【Pytorch】当num_workers ＞ 0时，程序卡住

问题描述：Dataloader设置num_workers > 0的时候，程序会在运行的某个时候卡住（随机的，不确定是什么时候）

解决方案：在Dataloader中，加入persistent_workers=True, worker_init_fn=worker_init_fn

def worker_init_fn(worker_id):
    worker_info = torch.utils.data.get_worker_info()
    dataset = worker_info.dataset
    if dataset.args.dataset == 'AVE':
        dataset.all_visual_pretrained_features = np.load(dataset.visual_pretrained_feature_path, allow_pickle=True).item()
    else:
        dataset.all_visual_pretrained_features = h5py.File(dataset.visual_pretrained_feature_path, 'r')
    dataset.all_audio_pretrained_features = np.load(dataset.audio_pretrained_feature_path, allow_pickle=True).item()


train_loader = DataLoader(train_data_set, batch_size=args.train_batch_size, num_workers=args.num_workers,
                              pin_memory=True, drop_last=True, shuffle=True, persistent_workers=True, worker_init_fn=worker_init_fn)