使用具有概率分布的 Pytorch Dataloader

Using Pytorch Dataloader with Probability Distribution

TL;DR:我想使用 DataLoader 对可用行进行 加权 随机抽样。怎么办?


我整理了一些 python 代码,这些代码通过对定义模型的参数使用批量梯度下降来拟合某种输入驱动动力系统和数据。我有以下 Python 代码片段,可以使用 Pytorch 完成工作。

k_trn = self.linear.k_gen(in_trn,t)
u_trn = torch.tensor(in_trn.T)
x_trn = torch.tensor(out_trn.T, dtype = torch.float)
data = TensorDataset(u_trn[:-1,:],k_trn[:-1,:],x_trn[1:,:])
loader = DataLoader(data, batch_size = 20, shuffle = True)

数据类型:

u_trn,k_trn,x_trn行对应三个轨迹(u对应“输入”)。每次我遍历加载器(这可以完成,例如使用循环 for u,k,x in loader:),我从 u_trn、20 行 k_trn 和 20 行中得到一批x_trn。这些行 select 以均匀概率编辑,没有替换。

要注意的是,我想以 非均匀 概率对这些行进行采样。特别地:表示 S = (1/1 + 1/2 + ... + 1/N)。我希望装载机 select 第 j 行的概率为 1/(S*j).

在查看 the relevant doumentation, 之后,我怀疑这可以通过在初始化 DataLoader 对象时弄乱 samplerbatch_sampler 关键字参数来完成,但是我我无法很好地解析文档以实现我正在寻找的行为。

如果您对此有任何帮助,我将不胜感激。我尽量让我的问题简短;如果我遗漏了任何相关信息,请告诉我。


跟进:在 Shai 的回答的帮助下,我已经让事情正常进行了。这是我用来测试它并确保一切都按预期工作的快速脚本。

将 numpy 导入为 np 进口手电筒 来自 torch.utils.data 导入 DataLoader、TensorDataset、WeightedRandomSampler 将 matplotlib.pyplot 导入为 plt

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler
import matplotlib.pyplot as plt

N = 100
x = np.zeros((N,2))
x[:,0] = 1 + np.arange(N)
data = TensorDataset(torch.Tensor(x))

weights = [1/j for j in range(1, N+1)]  # my weights
sampler = WeightedRandomSampler(weights, 10000, replacement=True)
loader = DataLoader(data, batch_size=20, sampler=sampler)
sums = []

for y, in loader:
    for k in range(len(y)):
        sums.append(np.sum(y[k].numpy()))

h = plt.hist(sums, bins = N)
a = h[0][0]
plt.plot([a/(n+1) for n in range(N)], lw = 3)

结果图:

注意权重是自动归一化的,所以不需要除以总和S。还要注意在loader中不需要shuffle=True;采样器自己负责随机化。

你为什么不简单地使用 WeightedRandomSampler

weights = [1./(S*j) for j in range(1, N+1)]  # your weights
sampler = WeightedRandomSampler(weights, replacement=True)
loader = DataLoader(data, batch_size=20, sampler=sampler)