def main():
model = UNet().to(device)
criterion = BinaryDiceLoss().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
train_dataset = ImageDataset(split='train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=train_batch_size,
num_workers=num_workers,
shuffle=True,
pin_memory=True)
valid_dataset = ImageDataset(split='valid')
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=valid_batch_size,
num_workers=num_workers,
shuffle=False,
pin_memory=True)
for epoch in range(epochs):
train(epoch, model, train_loader, criterion, optimizer)
valid(epoch, model, valid_loader, criterion)
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
},
save_file='pretrained/unet.pth.tar',
is_best=False)
def main():
global args, logger
args = parser.parse_args()
# logger = Logger(add_prefix(args.prefix, 'logs'))
set_prefix(args.prefix, __file__)
model = UNet(3, depth=5, in_channels=3)
print(model)
print('load unet with depth=5')
if args.cuda:
model = DataParallel(model).cuda()
else:
raise RuntimeError('there is no gpu')
criterion = nn.L1Loss(reduce=False).cuda()
print('use l1_loss')
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# accelerate the speed of training
cudnn.benchmark = True
data_loader = get_dataloader()
# class_names=['LESION', 'NORMAL']
# class_names = data_loader.dataset.class_names
# print(class_names)
since = time.time()
print('-' * 10)
for epoch in range(1, args.epochs + 1):
train(data_loader, model, optimizer, criterion, epoch)
if epoch % 40 == 0:
validate(model, epoch, data_loader)
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
validate(model, args.epochs, data_loader)
# save model parameter
torch.save(model.state_dict(),
add_prefix(args.prefix, 'identical_mapping.pkl'))
# save running parameter setting to json
write(vars(args), add_prefix(args.prefix, 'paras.txt'))