一、问题

使用了自己创建的dataset和collate_fn,使用了默认的dataloader,当设置num_worker的时候,超过0则报如下错误:

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

二、原因

查看了平时其他人的代码,发现自己在collate_fn对tensor使用了.to(device),而别人则是在从dataloader取出数据后,放入模型前,将其移到device中,例如:

    for epoch in range(max_epoch):
        for batch_idx, (en_id, att_mask_en, zh_id, att_mask_zh) in enumerate(data_loader): 
                en_id = en_id.to(device)

而我则是错误地在collate_fn中将数据移至device中:

def collate_fn(batch):
    en, zh = list(zip(*batch))
    en_output = tokenizer.encode_batch(en,is_pretokenized=True)
    zh_id,att_mask_zh = zh_mapper.encode_batch(zh)
    en_id =[]
    att_mask_en = []
    for item in en_output:
        en_id.append(item.ids)
        att_mask_en.append(item.attention_mask)
    en_id = torch.tensor(en_id, dtype=torch.long).to(device)
    att_mask_en = torch.tensor(att_mask_en, dtype=torch.bool).to(device)
    zh_id = torch.tensor(zh_id, dtype=torch.long).to(device)
    att_mask_zh = torch.tensor(att_mask_zh, dtype=torch.bool).to(device)

    return en_id, att_mask_en, zh_id, att_mask_zh

当然,在定义dataset,方法__getitem__ 里把张量放到device里也会产生同样的错误。

此外,我还使用了官方建议的spawn:

if __name__ == '__main__':
    torch.multiprocessing.set_start_method('spawn')
    train()

但是gpu会卡住不运行数据加载。 

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐