def train(model,
train_loader,
test_loader,
mode='EDSR_Baseline',
save_image_every=50,
save_model_every=10,
test_model_every=1,
epoch_start=0,
num_epochs=1000,
device=None,
refresh=True):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim,
step_size=1000,
gamma=0.99)
criterion = torch.nn.L1Loss()
GMSD = GMSD_quality().to(device)
mshf = MSHF(3, 3).to(device)
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
blurs = {}
for ksize in [3, 5, 7]:
blurs[ksize] = {}
for sigma in [0.4, 0.8, 1.0, 1.2, 1.6, 2.0]:
blurs[ksize][sigma] = Blur(ksize=ksize, sigma=sigma).to(device)
noise_sigma = 0.3
for epoch in range(epoch_start, epoch_start + num_epochs):
if epoch == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=
f'{mode} || Warming Up || Test Epoch {epoch}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
# hr = hr.to(device)
blur = blurs[7][2.0]
lr_input = blur(lr)
lr_input = lr_input + torch.rand_like(
lr, device=lr.device) * noise_sigma
_, features = model(lr_input)
dr = features[0]
# sr = quantize(sr)
psnr, ssim, msssim = evaluate(lr, dr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr'] = f'{psnr:.4f}'
pfix_test['ssim'] = f'{ssim:.4f}'
pfix_test['msssim'] = f'{msssim:.4f}'
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader,
desc=f'{mode} || Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
# hr = hr.to(device)
# prediction
ksize_ = random.choice([3, 5, 7])
sigma_ = random.choice([0.4, 0.8, 1.0, 1.2, 1.6, 2.0])
blur = blurs[ksize_][sigma_]
dnd = random.choice(['blur', 'noise', 'blur_and_noise'])
if dnd == 'blur':
lr_input = blur(lr)
elif dnd == 'noise':
lr_input = lr + torch.rand_like(
lr, device=lr.device) * noise_sigma
else:
lr_input = blur(lr)
lr_input = lr_input + torch.rand_like(
lr, device=lr.device) * noise_sigma
_, features = model(lr_input)
dr = features[0]
gmsd = GMSD(lr, dr)
# training
loss = criterion(lr, dr)
loss_tot = loss
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss'] = f'{loss.item():.4f}'
psnr, ssim, msssim = evaluate(lr, dr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix['PSNR'] = f'{psnr:.2f}'
pfix['SSIM'] = f'{ssim:.4f}'
# pfix['MSSSIM'] = f'{msssim:.4f}'
pfix['PSNR_mean'] = f'{psnr_mean:.2f}'
pfix['SSIM_mean'] = f'{ssim_mean:.4f}'
# pfix['MSSSIM_mean'] = f'{msssim_mean:.4f}'
free_gpu = get_gpu_memory()[0]
pfix['free GPU'] = f'{free_gpu}MiB'
pfix['Elapsed'] = f'{elapsed}'
pbar.set_postfix(pfix)
losses.append(loss.item())
if step % save_image_every == 0:
imsave([lr_input[0], dr[0], lr[0], gmsd[0]],
f'{result_dir}/epoch_{epoch+1}_iter_{step:05d}.jpg')
step += 1
logger.add_scalar("Loss/train", np.array(losses).mean(), epoch + 1)
logger.add_scalar("PSNR/train", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/train", ssim_mean, epoch + 1)
if (epoch + 1) % save_model_every == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch + 1) % test_model_every == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=f'{mode} || Test Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
fname = fname[0].split('/')[-1].split('.pt')[0]
lr = lr.to(device)
# hr = hr.to(device)
blur = blurs[7][2.0]
lr_input = blur(lr)
lr_input = lr_input + torch.rand_like(
lr, device=lr.device) * noise_sigma
_, features = model(lr_input)
dr = features[0]
mshf_lr = mshf(lr)
mshf_dr = mshf(dr)
gmsd = GMSD(lr, dr)
psnr, ssim, msssim = evaluate(lr, dr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr'] = f'{psnr:.4f}'
pfix_test['ssim'] = f'{ssim:.4f}'
pfix_test['msssim'] = f'{msssim:.4f}'
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
imsave([lr_input[0], dr[0], lr[0], gmsd[0]],
f'{result_dir}/{fname}.jpg')
mshf_vis = torch.cat(
(torch.cat([
mshf_dr[:, i, :, :]
for i in range(mshf_dr.shape[1])
],
dim=-1),
torch.cat([
mshf_lr[:, i, :, :]
for i in range(mshf_lr.shape[1])
],
dim=-1)),
dim=-2)
imsave(mshf_vis, f'{result_dir}/MSHF_{fname}.jpg')
hist['epoch'].append(epoch + 1)
hist['psnr'].append(psnr_mean)
hist['ssim'].append(ssim_mean)
hist['ms-ssim'].append(msssim_mean)
logger.add_scalar("PSNR/test", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/test", ssim_mean, epoch + 1)
logger.add_scalar("MS-SSIM/test", msssim_mean,
epoch + 1)
df = pd.DataFrame(hist)
df.to_csv(csv)
if args.n_epoch > 2:
argmax = t_list.argmax()
argmin = t_list.argmin()
t_list[argmax] = 0
t_list[argmin] = 0
t_avg = np.sum(t_list) / (args.n_epoch - 2)
else:
t_avg = np.sum(t_list) / args.n_epoch
print(
'Finish %d images for %d times in %.4fs, speed = %.4f image/s (%.4f ms/image)'
% (args.n_sample, args.n_epoch, t_end - t_start, args.n_sample / t_avg,
t_avg * 1000.0 / args.n_sample))
print('===================== benchmark finished =====================')
from utils import get_gpu_memory
gpu_mem = get_gpu_memory()
# save results
res_dir = 'cache/results'
if not os.path.exists(res_dir):
os.makedirs(res_dir)
res_file_path = os.path.join(
res_dir,
'%s_%s_%s_%d.txt' % (DLLIB, args.network, args.dtype, args.batch_size))
with open(res_file_path, 'w') as fd:
fd.write('%s %s %s %d %f %d\n' %
(DLLIB, args.network, args.dtype, args.batch_size,
args.n_sample / t_avg, gpu_mem))
def train(model,
train_loader,
test_loader,
mode='EDSR_Baseline',
save_image_every=50,
save_model_every=10,
test_model_every=1,
num_epochs=1000,
device=None,
refresh=True):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim,
step_size=1000,
gamma=0.99)
criterion = torch.nn.L1Loss()
######
ED = Edge().to(device)
######
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
for epoch in range(num_epochs):
if epoch == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=
f'Mode: {mode} || Warming Up || Test Epoch {epoch}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
sr, _, features = model(lr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr'] = f'{psnr:.4f}'
pfix_test['ssim'] = f'{ssim:.4f}'
pfix_test['msssim'] = f'{msssim:.4f}'
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader,
desc=f'Mode: {mode} || Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
hr = hr.to(device)
# prediction
sr, lr_edge, features = model(lr)
####
hr_edge = ED(hr)
loss_edge = criterion(lr_edge, hr_edge)
#####
# training
loss = criterion(hr, sr)
loss_tot = loss + 0.1 * loss_edge
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss'] = f'{loss.item():.4f}'
pfix['Loss Edge'] = f'{loss_edge.item():.4f}'
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix['PSNR'] = f'{psnr:.2f}'
pfix['SSIM'] = f'{ssim:.4f}'
# pfix['MSSSIM'] = f'{msssim:.4f}'
pfix['PSNR_mean'] = f'{psnr_mean:.2f}'
pfix['SSIM_mean'] = f'{ssim_mean:.4f}'
# pfix['MSSSIM_mean'] = f'{msssim_mean:.4f}'
free_gpu = get_gpu_memory()[0]
pfix['free GPU'] = f'{free_gpu}MiB'
pfix['Elapsed'] = f'{elapsed}'
pbar.set_postfix(pfix)
losses.append(loss.item())
if step % save_image_every == 0:
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0]],
f'{result_dir}/epoch_{epoch+1}_iter_{step:05d}.jpg')
step += 1
logger.add_scalar("Loss/train", np.array(losses).mean(), epoch + 1)
logger.add_scalar("PSNR/train", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/train", ssim_mean, epoch + 1)
if (epoch + 1) % save_model_every == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch + 1) % test_model_every == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=
f'Mode: {mode} || Test Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
fname = fname[0].split('/')[-1].split('.pt')[0]
lr = lr.to(device)
hr = hr.to(device)
sr, _, features = model(lr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr'] = f'{psnr:.4f}'
pfix_test['ssim'] = f'{ssim:.4f}'
pfix_test['msssim'] = f'{msssim:.4f}'
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0]],
f'{result_dir}/{fname}.jpg')
hist['epoch'].append(epoch + 1)
hist['psnr'].append(psnr_mean)
hist['ssim'].append(ssim_mean)
hist['ms-ssim'].append(msssim_mean)
logger.add_scalar("PSNR/test", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/test", ssim_mean, epoch + 1)
logger.add_scalar("MS-SSIM/test", msssim_mean,
epoch + 1)
df = pd.DataFrame(hist)
df.to_csv(csv)
def train(model,
train_loader,
test_loader,
mode='EDSR_Baseline',
save_image_every=50,
save_model_every=10,
test_model_every=1,
epoch_start=0,
num_epochs=1000,
device=None,
refresh=True,
scale=2,
today=None):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if today is None:
today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim,
step_size=1000,
gamma=0.99)
criterion = torch.nn.L1Loss()
GMSD = GMSD_quality().to(device)
opening = Opening().to(device)
blur = Blur().to(device)
mshf = MSHF(3, 3).to(device)
downx2_bicubic = nn.Upsample(scale_factor=1 / 2,
mode='bicubic',
align_corners=False)
downx4_bicubic = nn.Upsample(scale_factor=1 / 4,
mode='bicubic',
align_corners=False)
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
soft_mask = False
# hf_kernel = get_hf_kernel(mode='high')
for epoch in range(epoch_start, epoch_start + num_epochs):
if epoch == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=
f'{mode} || Warming Up || Test Epoch {epoch}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
sr, srx2, srx1 = model(lr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader,
desc=f'{mode} || Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
hr = hr.to(device)
hrx1 = downx4_bicubic(hr)
hrx2 = downx2_bicubic(hr)
# prediction
sr, srx2, srx1 = model(lr)
gmsd = GMSD(hr, sr)
sr_ = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr_)
if psnr >= 40 - 2 * scale:
soft_mask = True
else:
soft_mask = False
if soft_mask:
# with torch.no_grad():
# for _ in range(10): gmsd = opening(gmsd)
gmask = gmsd / gmsd.max()
gmask = (gmask > 0.2) * 1.0
gmask = blur(gmask)
gmask = (gmask - gmask.min()) / (gmask.max() -
gmask.min() + 1e-7)
gmask = (gmask + 0.25) / 1.25
gmask = gmask.detach()
gmaskx2 = downx2_bicubic(gmask)
gmaskx1 = downx4_bicubic(gmask)
# training
loss = criterion(sr * gmask, hr * gmask)
lossx2 = criterion(srx2 * gmaskx2, hrx2 * gmaskx2)
lossx1 = criterion(srx1 * gmaskx1, hrx1 * gmaskx1)
else:
loss = criterion(sr, hr)
lossx2 = criterion(srx2, hrx2)
lossx1 = criterion(srx1, hrx1)
# training
loss_tot = loss + 0.25 * lossx2 + 0.125 * lossx1
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Loss'] = f'{loss.item():.4f}'
pfix['x2'] = f'{lossx2.item():.4f}'
pfix['x1'] = f'{lossx1.item():.4f}'
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix['PSNR_mean'] = f'{psnr_mean:.2f}'
pfix['SSIM_mean'] = f'{ssim_mean:.4f}'
free_gpu = get_gpu_memory()[0]
pfix['Elapsed'] = f'{elapsed}'
pfix['free GPU'] = f'{free_gpu}MiB'
pbar.set_postfix(pfix)
losses.append(loss.item())
if step % save_image_every == 0:
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0], gmsd[0]],
f'{result_dir}/epoch_{epoch+1}_iter_{step:05d}.jpg')
step += 1
logger.add_scalar("Loss/train", np.array(losses).mean(), epoch + 1)
logger.add_scalar("PSNR/train", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/train", ssim_mean, epoch + 1)
if (epoch + 1) % save_model_every == 0:
torch.save(model.state_dict(),
f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch + 1) % test_model_every == 0:
with torch.no_grad():
with tqdm(
test_loader,
desc=f'{mode} || Test Epoch {epoch+1}/{num_epochs}',
position=0,
leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
fname = fname[0].split('/')[-1].split('.pt')[0]
lr = lr.to(device)
hr = hr.to(device)
sr, _, _ = model(lr)
mshf_hr = mshf(hr)
mshf_sr = mshf(sr)
gmsd = GMSD(hr, sr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
z = torch.zeros_like(lr[0])
_, _, llr, _ = lr.shape
_, _, hlr, _ = hr.shape
if hlr // 2 == llr:
xz = torch.cat((lr[0], z), dim=-2)
elif hlr // 4 == llr:
xz = torch.cat((lr[0], z, z, z), dim=-2)
imsave([xz, sr[0], hr[0], gmsd[0]],
f'{result_dir}/{fname}.jpg')
mshf_vis = torch.cat(
(torch.cat([
mshf_sr[:, i, :, :]
for i in range(mshf_sr.shape[1])
],
dim=-1),
torch.cat([
mshf_hr[:, i, :, :]
for i in range(mshf_hr.shape[1])
],
dim=-1)),
dim=-2)
imsave(mshf_vis, f'{result_dir}/MSHF_{fname}.jpg')
hist['epoch'].append(epoch + 1)
hist['psnr'].append(psnr_mean)
hist['ssim'].append(ssim_mean)
hist['ms-ssim'].append(msssim_mean)
logger.add_scalar("PSNR/test", psnr_mean, epoch + 1)
logger.add_scalar("SSIM/test", ssim_mean, epoch + 1)
logger.add_scalar("MS-SSIM/test", msssim_mean,
epoch + 1)
df = pd.DataFrame(hist)
df.to_csv(csv)
return model
position=0,
leave=True) as pbar:
for lr, hr, _ in pbar:
# prediction
pred = model(lr)
# training
loss = criterion(hr, pred)
optim.zero_grad()
loss.backward()
optim.step()
scheduler.step()
# training history
free_gpu = get_gpu_memory()[0]
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss'] = f'{loss.item():.4f}'
pfix['free GPU'] = f'{free_gpu}MiB'
pfix['Elapsed'] = f'{elapsed}'
hist['Iter'].append(step)
hist['Loss'].append(loss.item())
pbar.set_postfix(pfix)
if step % save_image_every == 0:
z = torch.zeros_like(lr[0])
xz = torch.cat((lr[0], z), dim=-2)
img = torch.cat((xz, pred[0], hr[0]), dim=-1)
parser.add_argument("-e","--epochs", type=int, default=25)
parser.add_argument("-f","--finetune", type=int, default=0)
args = parser.parse_args()
print("Arguments: {}".format(args))
model_names = available_models()
if args.model_name not in model_names:
raise ValueError("Wrong model name, please select one among {}".format(model_names))
print("OS:", sys.platform)
print("Python:", sys.version)
print("PyTorch:", torch.__version__)
print("Numpy:", np.__version__)
print("Number of CPU processors:", get_number_processors())
print("GPU:", get_gpu_name())
print("GPU memory: {}".format(get_gpu_memory()))
print("CUDA version:", get_cuda_version())
print("CUDNN version:", torch.backends.cudnn.version())
torch.backends.cudnn.benchmark=True # enables cudnn's auto-tuner
# Datasets
dataset = create_dataset(args.dataset_path, batch_size=args.batch_size)
# Training
if args.finetune:
model, metrics = finetune(dataset, args.model_name, SETS, args.epochs, args.gpus,
args.learning_rate, args.momentum, args.learning_rate_step,
args.learning_rate_epochs, verbose=True)
for content_id in content_ids:
new_result['content_ids'].append(content_id)
if content_id in content_ids_map and content_ids_map[
content_id] in result:
new_result['ratings'].append(
float(result[content_ids_map[content_id]]))
else:
new_result['ratings'].append(999.0) # TODO Fake data
return make_response(jsonify(new_result), STATUS_OK)
print("OS: ", sys.platform)
print("Python: ", sys.version)
print("PyTorch: ", torch.__version__)
print("Numpy: ", np.__version__)
print("Number of CPU processors: ", get_number_processors())
print("GPU: ", get_gpu_name())
print("GPU memory: ", get_gpu_memory())
print("CUDA: ", get_cuda_version())
print("USE_GPU: ", USE_GPU)
#Load data and model as global variables
data_layer, inv_userIdMap, inv_itemIdMap = load_train_data(TRAIN)
rencoder_api = load_recommender(data_layer.vector_dim, HIDDEN, ACTIVATION,
DROPOUT, MODEL_PATH)
muids_map, content_ids_map = load_train_muid_and_content_id(
TRAIN_MUID, TRAIN_CONTENT_ID)
if __name__ == "__main__":
run_server()
def train(model, model_sr, train_loader, test_loader, mode='EDSR_Baseline', save_image_every=50, save_model_every=10, test_model_every=1, epoch_start=0, num_epochs=1000, device=None, refresh=True):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
model_sr = model_sr.to(device)
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=1000, gamma= 0.99)
criterion = torch.nn.L1Loss()
GMSD = GMSD_quality().to(device)
mshf = MSHF(3, 3).to(device)
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
for epoch in range(epoch_start, epoch_start+num_epochs):
if epoch == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Warming Up || Test Epoch {epoch}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
sr, deep = model_sr(lr)
fake = model(sr)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr'] = f'{psnr:.4f}'
pfix_test['ssim'] = f'{ssim:.4f}'
pfix_test['msssim'] = f'{msssim:.4f}'
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader, desc=f'{mode} || Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
hr = hr.to(device)
# prediction
sr, deep = model_sr(lr)
fake = model(sr)
loss_fake = criterion(fake, torch.zeros_like(fake, device=fake.device))
real = model(hr)
loss_real = criterion(real, torch.ones_like(real, device=real.device))
# training
loss_tot = loss_fake + loss_real
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss real'] = f'{loss_real.item():.4f}'
pfix['Loss fake'] = f'{loss_fake.item():.4f}'
free_gpu = get_gpu_memory()[0]
pbar.set_postfix(pfix)
step += 1
if (epoch+1) % save_model_every == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch+1) % test_model_every == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Test Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
# prediction
sr, deep = model_sr(lr)
fake = model(sr)
loss_fake = criterion(fake, torch.zeros_like(fake, device=fake.device))
real = model(hr)
loss_real = criterion(real, torch.ones_like(real, device=real.device))
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix_test['Step'] = f'{step+1}'
pfix_test['Loss real'] = f'{loss_real.item():.4f}'
pfix_test['Loss fake'] = f'{loss_fake.item():.4f}'
pbar_test.set_postfix(pfix_test)
def train(model, train_loader, test_loader, mode='EDSR_Baseline', save_image_every=50, save_model_every=10, test_model_every=1, epoch_start=0, num_epochs=1000, device=None, refresh=True, scale=2, today=None):
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if today is None:
today = datetime.datetime.now().strftime('%Y.%m.%d')
result_dir = f'./results/{today}/{mode}'
weight_dir = f'./weights/{today}/{mode}'
logger_dir = f'./logger/{today}_{mode}'
csv = f'./hist_{today}_{mode}.csv'
if refresh:
try:
shutil.rmtree(result_dir)
shutil.rmtree(weight_dir)
shutil.rmtree(logger_dir)
except FileNotFoundError:
pass
os.makedirs(result_dir, exist_ok=True)
os.makedirs(weight_dir, exist_ok=True)
os.makedirs(logger_dir, exist_ok=True)
logger = SummaryWriter(log_dir=logger_dir, flush_secs=2)
model = model.to(device)
params = list(model.parameters())
optim = torch.optim.Adam(params, lr=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=1000, gamma= 0.99)
criterion = torch.nn.L1Loss()
start_time = time.time()
print(f'Training Start || Mode: {mode}')
step = 0
pfix = OrderedDict()
pfix_test = OrderedDict()
hist = dict()
hist['mode'] = f'{today}_{mode}'
for key in ['epoch', 'psnr', 'ssim', 'ms-ssim']:
hist[key] = []
soft_mask = False
sigmas = [10/255, 30/255, 50/255]
# hf_kernel = get_hf_kernel(mode='high')
for epoch in range(epoch_start, epoch_start+num_epochs):
sigma = 0.0004 * (epoch+1)
if epoch == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if epoch == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Warming Up || Test Epoch {epoch}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
lr = lr.to(device)
hr = hr.to(device)
sigma = np.random.choice(sigmas)
hr_input = hr + torch.randn_like(hr, device=hr.device)*sigma
hr_input = torch.clamp(hr_input, 0, 1)
sr = model(hr_input)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
if len(psnrs) > 1: break
with tqdm(train_loader, desc=f'{mode} || Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar:
psnrs = []
ssims = []
msssims = []
losses = []
for lr, hr, _ in pbar:
lr = lr.to(device)
hr = hr.to(device)
sigma = np.random.choice(sigmas)
hr_input = hr + torch.randn_like(hr, device=hr.device)*sigma
hr_input = torch.clamp(hr_input, 0, 1)
sr = model(hr_input)
sr_ = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr_)
loss = criterion(sr, hr)
# training
loss_tot = loss
optim.zero_grad()
loss_tot.backward()
optim.step()
scheduler.step()
# training history
elapsed_time = time.time() - start_time
elapsed = sec2time(elapsed_time)
pfix['Step'] = f'{step+1}'
pfix['Loss'] = f'{loss.item():.4f}'
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix['PSNR_mean'] = f'{psnr_mean:.2f}'
pfix['SSIM_mean'] = f'{ssim_mean:.4f}'
free_gpu = get_gpu_memory()[0]
pfix['free GPU'] = f'{free_gpu}MiB'
pfix['Elapsed'] = f'{elapsed}'
pbar.set_postfix(pfix)
losses.append(loss.item())
step += 1
logger.add_scalar("Loss/train", np.array(losses).mean(), epoch+1)
logger.add_scalar("PSNR/train", psnr_mean, epoch+1)
logger.add_scalar("SSIM/train", ssim_mean, epoch+1)
if (epoch+1) % save_model_every == 0:
torch.save(model.state_dict(), f'{weight_dir}/epoch_{epoch+1:04d}.pth')
if (epoch+1) % test_model_every == 0:
with torch.no_grad():
with tqdm(test_loader, desc=f'{mode} || Test Epoch {epoch+1}/{num_epochs}', position=0, leave=True) as pbar_test:
psnrs = []
ssims = []
msssims = []
for lr, hr, fname in pbar_test:
fname = fname[0].split('/')[-1].split('.pt')[0]
# lr = lr.to(device)
hr = hr.to(device)
sigma = np.random.choice(sigmas)
hr_input = hr + torch.randn_like(hr, device=hr.device)*sigma
hr_input = torch.clamp(hr_input, 0, 1)
sr = model(hr_input)
sr = quantize(sr)
psnr, ssim, msssim = evaluate(hr, sr)
psnrs.append(psnr)
ssims.append(ssim)
msssims.append(msssim)
psnr_mean = np.array(psnrs).mean()
ssim_mean = np.array(ssims).mean()
msssim_mean = np.array(msssims).mean()
pfix_test['psnr_mean'] = f'{psnr_mean:.4f}'
pfix_test['ssim_mean'] = f'{ssim_mean:.4f}'
pfix_test['msssim_mean'] = f'{msssim_mean:.4f}'
pbar_test.set_postfix(pfix_test)
imsave([hr_input[0], sr[0], hr[0]], f'{result_dir}/{fname}.jpg')
hist['epoch'].append(epoch+1)
hist['psnr'].append(psnr_mean)
hist['ssim'].append(ssim_mean)
hist['ms-ssim'].append(msssim_mean)
logger.add_scalar("PSNR/test", psnr_mean, epoch+1)
logger.add_scalar("SSIM/test", ssim_mean, epoch+1)
logger.add_scalar("MS-SSIM/test", msssim_mean, epoch+1)
df = pd.DataFrame(hist)
df.to_csv(csv)
