train_dir = opt.TRAINING.TRAIN_DIR
val_dir = opt.TRAINING.VAL_DIR
save_images = opt.TRAINING.SAVE_IMAGES
######### Model ###########
model_restoration = MIRNet()
model_restoration.cuda()
device_ids = [i for i in range(torch.cuda.device_count())]
if torch.cuda.device_count() > 1:
print("\n\nLet's use", torch.cuda.device_count(), "GPUs!\n\n")
new_lr = opt.OPTIM.LR_INITIAL
optimizer = optim.Adam(model_restoration.parameters(), lr=new_lr, betas=(0.9, 0.999),eps=1e-8, weight_decay=1e-8)
######### Scheduler ###########
if warmup:
warmup_epochs = 3
scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, opt.OPTIM.NUM_EPOCHS-warmup_epochs, eta_min=1e-6)
scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=warmup_epochs, after_scheduler=scheduler_cosine)
scheduler.step()
######### Resume ###########
if opt.TRAINING.RESUME:
path_chk_rest = utils.get_last_path(model_dir, '_latest.pth')
utils.load_checkpoint(model_restoration,path_chk_rest)
start_epoch = utils.load_start_epoch(path_chk_rest) + 1
utils.load_optim(optimizer, path_chk_rest)
train_dir = opt.TRAINING.TRAIN_DIR
val_dir = opt.TRAINING.VAL_DIR
save_images = opt.TRAINING.SAVE_IMAGES
######### Model ###########
model_restoration = MIRNet()
model_restoration.cuda()
device_ids = [i for i in range(torch.cuda.device_count())]
if torch.cuda.device_count() > 1:
print("\n\nLet's use", torch.cuda.device_count(), "GPUs!\n\n")
new_lr = opt.OPTIM.LR_INITIAL
optimizer = optim.Adam(model_restoration.parameters(),
lr=new_lr,
betas=(0.9, 0.999),
eps=1e-8,
weight_decay=1e-8)
######### Resume ###########
if opt.TRAINING.RESUME:
path_chk_rest = utils.get_last_path(model_dir, '_latest.pth')
utils.load_checkpoint(model_restoration, path_chk_rest)
start_epoch = utils.load_start_epoch(path_chk_rest) + 1
lr = utils.load_optim(optimizer, path_chk_rest)
for p in optimizer.param_groups:
p['lr'] = lr
warmup = False
