PyTorch 动态更新 DataLoader

举个例子，我们只想取 Dataset 中的一部分，所以可以使用 SubsetRandomSampler 。

from torch.utils.data import DataLoader, TensorDataset, SubsetRandomSampler, BatchSampler
import torch

candidate = [1]
dataset = TensorDataset(torch.tensor(list(range(10))))
dataloader = DataLoader(dataset, batch_size=1, sampler=SubsetRandomSampler(candidate))

for idx, data in enumerate(dataloader):
    print(data)

输出的结果肯定是只取第二个数据，也就是

[tensor([1])]

一个简单的想法是，我们能否直接更新 dataloader.sampler 以及 dataloader.batch_sampler ？

candidate = [2]
dataloader.sampler = SubsetRandomSampler(candidate)
dataloader.batch_sampler = BatchSampler(SubsetRandomSampler(candidate), batch_size=1, drop_last=True)

遗憾的是，PyTorch 并不支持这个操作。

方法：
In-place 地改变 candidate 这个变量就行了。

from torch.utils.data import DataLoader, TensorDataset
import torch

candidate = [1,2]
dataset = TensorDataset(torch.tensor(list(range(10))))
dataloader = DataLoader(dataset, batch_size=1, sampler=SubsetRandomSampler(candidate))

for idx, data in enumerate(dataloader):
    print(data)

[candidate.pop() for i in range(len(candidate))]
candidate.extend([3,4])

for idx, data in enumerate(dataloader):
    print(data)

我们来看下效果：

[tensor([1])]
[tensor([2])]
[tensor([4])]
[tensor([3])]

可以看到这个方法是成功的。

中间两行：

[candidate.pop() for i in range(len(candidate))]
candidate.extend([3,4])

是先 in-place 地把这个 list 清空，然后再 in-place 地更新这个 list。

注意，这里的关键是 in-place！

如果你直接对 candidate 进行赋值，是无法达到修改效果的。

from torch.utils.data import DataLoader, TensorDataset
import torch

candidate = [1,2]
dataset = TensorDataset(torch.tensor(list(range(10))))
dataloader = DataLoader(dataset, batch_size=1, sampler=SubsetRandomSampler(candidate))

for idx, data in enumerate(dataloader):
    print(data)

candidate = [3,4]

for idx, data in enumerate(dataloader):
    print(data)

无法达到效果：

[tensor([1])]
[tensor([2])]
[tensor([2])]
[tensor([1])]