weight_decay = 5e-5
lr_schedule = 0.985
def adjust_lr(optimizer, current_lr, schedule):
current_lr = current_lr * schedule
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
return current_lr
if __name__ == "__main__":
args, unparsed = config.get_args()
model = Unet(args)
model = model.cuda()
model.train()
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
loss = MulticlassDiceLoss()
train = get_file_list(brats_preprocessed_folder, train_ids_path)
val = get_file_list(brats_preprocessed_folder, valid_ids_path)
shapes = [brats_dataloader.load_patient(i)[0].shape[1:] for i in train]
max_shape = np.max(shapes, 0)
max_shape = list(np.max((max_shape, patch_size), 0))
dataloader_train = brats_dataloader(train,
batch_size,
max_shape,
def train(args):
# device
device = torch.device("cuda:%d" %
args.gpu if torch.cuda.is_available() else "cpu")
# data
trainset = SonyDataset(args.input_dir, args.gt_dir, args.ps)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
logging.info("data loading okay")
# model
model = Unet().cuda()
print(model)
# resume
starting_epoch = 0
if args.resume is not None:
model.load_state_dict(torch.load(args.resume))
starting_epoch = int(args.resume[-7:-3])
print('resume at %d epoch' % starting_epoch)
# loss function
color_loss = nn.L1Loss()
gradient_loss = GradLoss(device)
# optimizer
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.wd)
# lr scheduler
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.1)
model.train()
# training
running_loss = 0.0
for epoch in range(starting_epoch + 1, starting_epoch + args.num_epoch):
scheduler.step()
for i, databatch in enumerate(train_loader):
input_patch, gt_patch, train_id, ratio = databatch
input_patch = input_patch.cuda()
gt_patch = gt_patch.cuda(async=True)
input_patch = torch.autograd.Variable(input_patch)
gt_patch = torch.autograd.Variable(gt_patch)
# # get the inputs
# input_patch, gt_patch, train_id, ratio = databatch
# input_patch, gt_patch = input_patch.to(device), gt_patch.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(input_patch)
loss = color_loss(outputs, gt_patch) + gradient_loss(
outputs, gt_patch)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % args.log_interval == (args.log_interval - 1):
print('[%d, %5d] loss: %.3f %s' %
(epoch, i, running_loss / args.log_interval,
datetime.now()))
running_loss = 0.0
if epoch % args.save_freq == 0:
if not os.path.isdir(
os.path.join(args.result_dir, '%04d' % epoch)):
os.makedirs(os.path.join(args.result_dir, '%04d' % epoch))
gt_patch = gt_patch.cpu().detach().numpy()
outputs = outputs.cpu().detach().numpy()
train_id = train_id.numpy()
ratio = ratio.numpy()
temp = np.concatenate(
(gt_patch[0, :, :, :], outputs[0, :, :, :]), axis=2)
scipy.misc.toimage(
temp * 255, high=255, low=0, cmin=0,
cmax=255).save(args.result_dir +
'%04d/%05d_00_train_%d.jpg' %
(epoch, train_id[0], ratio[0]))
# at the end of epoch
if epoch % args.model_save_freq == 0:
torch.save(model.state_dict(),
args.checkpoint_dir + './model_%d.pl' % epoch)