pytorch中使用多顯卡訓練以及訓練時報錯：expect more than 1 value per channel when training, got input size..

pytorch在訓練中使用多卡：

        conf.device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
        conf.device_ids = list(conf.device_ids)
        self.model = torch.nn.DataParallel(self.model, device_ids=conf.device_ids)
        self.model.to(conf.device)

 然后在訓練的命令行需要指定GPU：

export CUDA_VISIBLE_DEVICES=4,5,6,7

 在pytorch內部，conf.device_ids依舊是從0開始的；

 

訓練的時候報錯：

 

 

是由於batchnorm層需要大於一個樣本去計算其中的參數，網上搜索了解決方法是將dataloader的一個丟棄參數設置為true：

 

但是我設置后依舊會報錯，然后就在train里面人為修改了一下：

如果剩下的照片為1，那么就丟掉，就可以了：