def _init_model(self): M = Models() model = M.FPN(img_ch=3, output_ch=1) # model = U_Net(img_ch=3, output_ch=1) if torch.cuda.device_count() > 1 and self.args.mgpu: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = nn.DataParallel(model) self.model = model.to(self.device)
def _init_dataset(self): M = Models() if self.args.mgpu: self.batch_size = 28 print('batch_size: ', self.batch_size) self.date = '/2020-05-06~11:38:23' self.Mo = M.FPN(img_ch=3, output_ch=1) else: self.batch_size = 7 print('batch_size: ', self.batch_size) self.date = '/2020-05-25~05:51:58' self.Mo = U_Net(img_ch=3, output_ch=1) test_images = Angioectasias(self.abnormality, mode='test') self.test_queue = DataLoader(test_images, batch_size=self.batch_size, drop_last=False)