Блог

Я пытаюсь использовать несколько графических процессоров для кода, который уже работает корректно для одного графического процессора. Я внес эти изменения, чтобы сделать его способным разделять несколько графических процессоров:

 classifier = AnnotationCNN.load_from_checkpoint(classifier_ckpt,normalizer_ckpt=normalizer_ckpt)
classifier=classifier.eval()
    
**device = torch.device("cuda" if torch.cuda.is_available() else "cpu")**
classifier = classifier.to(torch.float16).to(device)
**classifier= nn.DataParallel(classifier)**


dataset = InitImageTifDataset(
                data_root=data_root,
                level=level,
                region_width=region_width,
                region_height=region_height,
                crop_size=crop_size,
                image_dir=image_dir,
                image_name=image_name,
                channels=channels
            )

        dl = DataLoader(dataset, batch_size=5, shuffle=False, num_workers=5)
        for i,(img,mean_stddev,filename) in enumerate(dl):
            mask = mean_stddev > 20

            img=img[torch.nonzero(mask)]
            if img.nelement() != 0 :
                **tensor=img[:,0,0,:,:,:].to(device).to(torch.float16).to(device)**
                *_,preds = classifier(tensor)*
 

Но у меня следующая ошибка:

 RuntimeError: Expected tensor for argument #1 'output' to have the same device as tensor for argument #5 'weight'; but device 0 does not equal 1 (while checking arguments for batch_norm_cuda)
 

Я почти уверен, что ошибка исходит от загрузчика данных, который не адаптирован к нескольким графическим процессорам. Как я могу это исправить, спасибо.