使用带有概率分布的 Pytorch 数据加载器

Question

TL;DR：我想使用DataLoader 对可用行进行加权随机抽样。怎么做？

---

我已经将一些 python 代码放在一起，这些代码使用批量梯度下降对定义模型的参数进行拟合，以将某种输入驱动的动态系统与数据相匹配。我有以下 Python 代码 sn-p 可以使用 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: N x 1 张量（pytorch 的数组）
• k_trn: N x K 张量
• x_trn: N x n 张量

u_trn,k_trn,x_trn 的行对应三个轨迹（其中 u 对应“输入”）。每次我迭代加载程序（可以完成，例如使用循环for u,k,x in loader:），我从 u_trn、20 行 k_trn 和 20 行 x_trn 中获得一批 20 行。这些行以均匀的概率被选择，无需替换。

问题是我想以非均匀概率对这些行进行抽样。特别是：表示 S = (1/1 + 1/2 + ... + 1/N)。我希望加载器以 1/(S*j) 的概率选择第 j 行。

查看the relevant doumentation, 后，我怀​​疑这可以通过在初始化DataLoader 对象时弄乱sampler 或batch_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。还要注意，加载器中不需要shuffle=True；采样器自己负责随机化。

Answer 1

你为什么不直接使用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)