问题描述: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)