PyTorch 动态更新 Dataset 对 DataLoader不生效

项目场景

使用 Stitcher 增强小目标检测效果。在每个 iter 训练结束后，根据小目标贡献的 loss 占总 loss 的比重，决定下一个 iter 的数据是否需要 stitch 。

环境：PyTorch 1.3

问题描述

实现上需要使用 Hook 统计 loss 占比，然后设置 Dataset 的标志位。 Dataset.__getitem__ 会首先查询标志位，决定是否返回 stitch 后的结果。运行时发现标志位设置失败， stitch 一直没有激活。

from mmcv.runner import BaseRunner, Hook, HOOKS
from mmdet.datasets.builder import DATASETS


@HOOKS.register_module()
class SmallFeedbackHook(Hook):
    def before_train_epoch(self, runner: BaseRunner):
        data_loader: DataLoader = runner.data_loader
        dataset: StitchedDataset = data_loader.dataset
        dataset.active = False
        self.dataset = dataset

    def after_train_iter(self, runner: BaseRunner):
        ratio = ...
        self.dataset.active = ratio <= self.thresh


@DATASETS.register_module()
class StitchedDataset:
    def __getitem__(self, idx: int) -> Dict[str, Any]:
        if self.active: ...
        else: ...

原因分析

在 Hook 中设置断点发现标志位设置成功，但是 Dataset.__getitem__ 打印的信息却显示标志位设置失败，于是猜测 DataLoader 在创建进程时复制了 Dataset 对象。

参考 ptrblck 提供的代码，可以证明 DataLoader 在创建进程时没有复制 Dataset 对象，而是使用了 Dataset 对象的引用

import os

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


class MyDataset(Dataset):
    LEN = 30
    
    def __init__(self):
        self.data = torch.zeros(self.LEN, 1)

    def __getitem__(self, index):
        print("Get item", index, "of", id(self), "at", os.getpid())
        return self.data[index]

    def __len__(self):
        return self.LEN

    def regenerate(self):
        print("Regenerate", id(self), "at", os.getpid())
        self.data = torch.ones(self.LEN, 1)


if __name__ == '__main__':
    dataset = MyDataset()
    data_loader = DataLoader(dataset, num_workers=3)

    print(dataset[0])  # output 0
    for d in data_loader:
        print(d)

    dataset.regenerate()

    print(dataset[0])  # output 1
    for d in data_loader:
        print(d)

输出

Get item 1 of 140308728889528 at 25433
tensor([0.])
Get item 1 of 140308728889528 at 31734
Get item 0 of 140308728889528 at 31727
Get item 2 of 140308728889528 at 31744
tensor([[0.]])
...
Regenerate 140308728889528 at 25433
Get item 0 of 140308728889528 at 25433
tensor([1.])
Get item 0 of 140308728889528 at 32116
Get item 1 of 140308728889528 at 32117
Get item 2 of 140308728889528 at 32118
tensor([[1.]])
...

可以看出， 25433 是主进程， 31727/31734/31744 是第一次遍历 DataLoader 创建的三个子进程， 32116/32117/32118是第二次遍历 DataLoader 创建的三个子进程。所有进程都使用 140308728889528 的 MyDataset 对象， MyDataset.regenerate 也生效了。但是如果在一次遍历的过程中调用 MyDataset.regenerate 就无法生效。

if __name__ == '__main__':
    dataset = MyDataset()
    data_loader = DataLoader(dataset, num_workers=3)

    for d in data_loader:
        print(d)
        dataset.regenerate()

输出

Get item 1 of 139961664745768 at 12979
Get item 2 of 139961664745768 at 12982
Get item 0 of 139961664745768 at 12977
Get item 5 of 139961664745768 at 12982
Get item 3 of 139961664745768 at 12977
Get item 4 of 139961664745768 at 12979
Get item 6 of 139961664745768 at 12977
tensor([[0.]])
Regenerate 139961664745768 at 9793
Get item 7 of 139961664745768 at 12979
tensor([[0.]])
Regenerate 139961664745768 at 9793
Get item 8 of 139961664745768 at 12982
tensor([[0.]])
...

目前猜测，如果对 Dataset 的更改发生在子进程创建之前，更改就能生效，否则不生效。

解决方案

目前没有较好的解决方案，只能将 DataLoader.__init__ 的 num_workers 设置为 0 。这样可以在主进程中读取数据，避免创建子进程，但同时会减慢训练速度。