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)
