Spaces:

yangheng
/

Super-Resolution-Anime-Diffusion

Running

App Files Files Community

Super-Resolution-Anime-Diffusion / Waifu2x /train.py

yangheng

init

9842c28 over 1 year ago

raw history blame contribute delete

No virus

6.05 kB

	from torch import optim
	from torch.utils.data import DataLoader
	from torchvision.utils import save_image
	from tqdm import trange

	from Dataloader import *
	from .utils import image_quality
	from .utils.cls import CyclicLR
	from .utils.prepare_images import *

	train_folder = "./dataset/train"
	test_folder = "./dataset/test"

	img_dataset = ImageDBData(
	db_file="dataset/images.db",
	db_table="train_images_size_128_noise_1_rgb",
	max_images=24,
	)
	img_data = DataLoader(img_dataset, batch_size=6, shuffle=True, num_workers=6)

	total_batch = len(img_data)
	print(len(img_dataset))

	test_dataset = ImageDBData(
	db_file="dataset/test2.db",
	db_table="test_images_size_128_noise_1_rgb",
	max_images=None,
	)
	num_test = len(test_dataset)
	test_data = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=1)

	criteria = nn.L1Loss()

	model = CARN_V2(
	color_channels=3,
	mid_channels=64,
	conv=nn.Conv2d,
	single_conv_size=3,
	single_conv_group=1,
	scale=2,
	activation=nn.LeakyReLU(0.1),
	SEBlock=True,
	repeat_blocks=3,
	atrous=(1, 1, 1),
	)

	model.total_parameters()


	# model.initialize_weights_xavier_uniform()

	# fp16 training is available in GPU only
	model = network_to_half(model)
	model = model.cuda()
	model.load_state_dict(torch.load("CARN_model_checkpoint.pt"))

	learning_rate = 1e-4
	weight_decay = 1e-6
	optimizer = optim.Adam(
	model.parameters(), lr=learning_rate, weight_decay=weight_decay, amsgrad=True
	)
	# optimizer = optim.SGD(model.parameters(), momentum=0.9, nesterov=True, weight_decay=weight_decay, lr=learning_rate)

	# optimizer = FP16_Optimizer(optimizer, static_loss_scale=128.0, verbose=False)
	# optimizer.load_state_dict(torch.load("CARN_adam_checkpoint.pt"))

	last_iter = -1 # torch.load("CARN_scheduler_last_iter")
	scheduler = CyclicLR(
	optimizer,
	base_lr=1e-4,
	max_lr=1e-4,
	step_size=3 * total_batch,
	mode="triangular",
	last_batch_iteration=last_iter,
	)
	train_loss = []
	train_ssim = []
	train_psnr = []

	test_loss = []
	test_ssim = []
	test_psnr = []

	# train_loss = torch.load("train_loss.pt")
	# train_ssim = torch.load("train_ssim.pt")
	# train_psnr = torch.load("train_psnr.pt")
	#
	# test_loss = torch.load("test_loss.pt")
	# test_ssim = torch.load("test_ssim.pt")
	# test_psnr = torch.load("test_psnr.pt")


	counter = 0
	iteration = 2
	ibar = trange(
	iteration,
	ascii=True,
	maxinterval=1,
	postfix={"avg_loss": 0, "train_ssim": 0, "test_ssim": 0},
	)
	for i in ibar:
	# batch_loss = []
	# insample_ssim = []
	# insample_psnr = []
	for index, batch in enumerate(img_data):
	scheduler.batch_step()
	lr_img, hr_img = batch
	lr_img = lr_img.cuda().half()
	hr_img = hr_img.cuda()

	# model.zero_grad()
	optimizer.zero_grad()
	outputs = model.forward(lr_img)
	outputs = outputs.float()
	loss = criteria(outputs, hr_img)
	# loss.backward()
	optimizer.backward(loss)
	# nn.utils.clip_grad_norm_(model.parameters(), 5)
	optimizer.step()

	counter += 1
	# train_loss.append(loss.item())

	ssim = image_quality.msssim(outputs, hr_img).item()
	psnr = image_quality.psnr(outputs, hr_img).item()

	ibar.set_postfix(
	ratio=index / total_batch,
	loss=loss.item(),
	ssim=ssim,
	batch=index,
	psnr=psnr,
	lr=scheduler.current_lr,
	)
	train_loss.append(loss.item())
	train_ssim.append(ssim)
	train_psnr.append(psnr)

	# +++++++++++++++++++++++++++++++++++++
	# save checkpoints by iterations
	# -------------------------------------

	if (counter + 1) % 500 == 0:
	torch.save(model.state_dict(), "CARN_model_checkpoint.pt")
	torch.save(optimizer.state_dict(), "CARN_adam_checkpoint.pt")
	torch.save(train_loss, "train_loss.pt")
	torch.save(train_ssim, "train_ssim.pt")
	torch.save(train_psnr, "train_psnr.pt")
	torch.save(scheduler.last_batch_iteration, "CARN_scheduler_last_iter.pt")

	# +++++++++++++++++++++++++++++++++++++
	# End of One Epoch
	# -------------------------------------

	# one_ite_loss = np.mean(batch_loss)
	# one_ite_ssim = np.mean(insample_ssim)
	# one_ite_psnr = np.mean(insample_psnr)

	# print(f"One iteration loss {one_ite_loss}, ssim {one_ite_ssim}, psnr {one_ite_psnr}")
	# train_loss.append(one_ite_loss)
	# train_ssim.append(one_ite_ssim)
	# train_psnr.append(one_ite_psnr)

	torch.save(model.state_dict(), "CARN_model_checkpoint.pt")
	# torch.save(scheduler, "CARN_scheduler_optim.pt")
	torch.save(optimizer.state_dict(), "CARN_adam_checkpoint.pt")
	torch.save(train_loss, "train_loss.pt")
	torch.save(train_ssim, "train_ssim.pt")
	torch.save(train_psnr, "train_psnr.pt")
	# torch.save(scheduler.last_batch_iteration, "CARN_scheduler_last_iter.pt")

	# +++++++++++++++++++++++++++++++++++++
	# Test
	# -------------------------------------

	with torch.no_grad():
	ssim = []
	batch_loss = []
	psnr = []
	for index, test_batch in enumerate(test_data):
	lr_img, hr_img = test_batch
	lr_img = lr_img.cuda()
	hr_img = hr_img.cuda()

	lr_img_up = model(lr_img)
	lr_img_up = lr_img_up.float()
	loss = criteria(lr_img_up, hr_img)

	save_image([lr_img_up[0], hr_img[0]], f"check_test_imgs/{index}.png")
	batch_loss.append(loss.item())
	ssim.append(image_quality.msssim(lr_img_up, hr_img).item())
	psnr.append(image_quality.psnr(lr_img_up, hr_img).item())

	test_ssim.append(np.mean(ssim))
	test_loss.append(np.mean(batch_loss))
	test_psnr.append(np.mean(psnr))

	torch.save(test_loss, "test_loss.pt")
	torch.save(test_ssim, "test_ssim.pt")
	torch.save(test_psnr, "test_psnr.pt")

	# import subprocess

	# subprocess.call(["shutdown", "/s"])