generic_train_test.py

import os
import time
from tqdm import tqdm
from datetime import datetime
import wandb

import torch
from torchvision.utils import make_grid, save_image

from utils import *
from utils.metrics import *
# from skimage.measure import compare_psnr
from data import getLoader

Train_test = False

class Generic_train_test():
	def __init__(self, opts, accelerator, net, optimizer, scheduler, train_loader, val_loaders, datasets_name, metrics):
		self.opts = opts
		self.net = net
		self.optimizer = optimizer
		self.scheduler = scheduler

		self.l1_loss = torch.nn.L1Loss()
		self.stop_l1 = opts['train']['stop_l1'] if 'stop_l1' in opts['train'] else opts['train']['epochs']
		self.ms_ssim = MS_SSIM(accelerator)
		# self.ssim = SSIM()
		self.loss_funs = opts['train']['loss_funs']
		self.loss_weights = opts['train']['loss_weights']

		self.train_loader = train_loader
		self.val_loaders = val_loaders
		self.datasets_name = datasets_name

		# metrics
		self.best_loss = metrics['val_loss']
		self.best_ssim = metrics['val_ssim']
		self.best_psnr = metrics['val_psnr']

		# dirs
		self.checkpoint_dir = opts['Experiment']['checkpoint_dir']
		self.result_dir = opts['Experiment']['result_dir']
		self.lambda_gray = 0.5

		self.sar_trans = opts['sar_trans'] if 'sar_trans' in opts else False
		self.use_id = opts['use_id'] if 'use_id' in opts else False
		self.change_dataset = opts['train']['change_dataset'] if 'change_dataset' in opts['train'] else None

	def decode_input(self, data):
		return data
		# raise NotImplementedError()

	def train(self, accelerator, run, start_epoch, end_epoch):
		if accelerator.is_local_main_process:
			wandb.watch(self.net)
			wandb.define_metric("epoch")
			wandb.define_metric("lr", step_metric="epoch")
			metrics = ['train_loss', 'train_ssim', 'train_psnr', 'val_loss', 'val_ssim', 'val_psnr']
			for metric in metrics:
				wandb.define_metric(metric, step_metric="epoch")

		accelerator.print(f"#Train dataset: {self.datasets_name[0]}")
		accelerator.print('#Train image nums: ', len(self.train_loader)*self.opts['datasets']['train']['batch_size'])
		for epoch in range(start_epoch+1, end_epoch+1):
			cureent_epoch = f'epoch_{epoch}'
			if self.change_dataset!=None and cureent_epoch in self.change_dataset.keys():
				change_opt = self.change_dataset[cureent_epoch]
				self.train_loader = accelerator.prepare(getLoader(change_opt))
				accelerator.print('#Change dataet and image nums: ',
								  len(self.train_loader) * self.opts['datasets']['train']['batch_size'])

			batch_time = AverageMeter('Time', ':6.3f')
			data_time = AverageMeter('Data', ':6.3f')
			m_l1_loss = AverageMeter('Loss', ':.4e')
			m_ssim = AverageMeter('SSIM', ':6.2f')
			m_psnr = AverageMeter('PSNR', ':6.2f')

			if accelerator.is_local_main_process:
				wandb.log({'lr': self.optimizer.param_groups[0]["lr"], 'epoch':epoch})
			self.net.train()
			end = time.time()
			with tqdm(total=len(self.train_loader), desc=f'[Epoch {epoch}/{end_epoch}]', unit='batch',
					  disable=not accelerator.is_local_main_process) as train_pbar:
				for step, batch in enumerate(self.train_loader):
					with accelerator.accumulate(self.net):
						data_time.update(time.time() - end)

						image = batch['opt_cloudy']
						sar = batch['sar']
						label = batch['opt_clear']
						if self.use_id:
							image_id = batch['image_id']
						self.optimizer.zero_grad()
						loss_all = 0
						if self.use_id:
							pred = self.net(image, sar, image_id, accelerator)
						elif self.sar_trans:
							pred = self.net(sar)
						else:
							pred = self.net(image, sar)

						if 'pixel' in self.loss_funs.keys() and epoch < self.stop_l1:
							loss_l1 = self.l1_loss(pred, label)
							loss_all += loss_l1 * self.loss_weights[0]
						if 'ssim' in self.loss_funs.keys():
							if self.loss_funs['ssim'] == 'ms_ssim':
								loss_ssim = 1 - self.ms_ssim(pred, label)
							else:
								loss_ssim = 1 - SSIM(pred, label)
							loss_all += loss_ssim * self.loss_weights[1]

						accelerator.backward(loss_all)

						self.optimizer.step()
						# self.scheduler.step()

						# metrics
						ssim = SSIM(pred, label).item() # ?
						psnr = PSNR(pred, label)

						# loss_v = torch.mean(accelerator.gather_for_metrics(loss)).item()
						m_l1_loss.update(loss_l1.item(), image.size(0)) # average ?
						m_ssim.update(ssim, image.size(0))
						m_psnr.update(psnr, image.size(0))

						batch_time.update(time.time() - end)
						end = time.time()
						if accelerator.is_local_main_process:
							# =========== visualize results ============#
							if step % self.opts['log_step_freq'] == 0:
								total_steps = len(self.train_loader) * (epoch-1) + step + 1
								wandb.log({'loss': m_l1_loss.avg, 'ssim': m_ssim.avg, 'psnr': m_psnr.avg, 'step':total_steps})

							if step % self.opts['visual_step_freq']==0 or step==len(self.train_loader)-1:
								# figure
								img_sample = torch.cat([image.data, pred.data, label.data], -1)  # 按宽拼接
								grid = make_grid(img_sample, nrow=1, normalize=True) # 每一行显示的图像列数
								save_image(grid, os.path.join(self.result_dir, 'train_images', f'img_epoch_{epoch}_step_{step}.png'))

							# 后缀信息
							train_pbar.set_postfix(ordered_dict={'loss': m_l1_loss.avg, 'ssim': m_ssim.avg, 'psnr': m_psnr.avg})
							train_pbar.update()

						if Train_test:
							break

			# if epoch > self.opts['train']['scheduler']['lr_start_epoch_decay'] - self.opts['train']['scheduler']['lr_step']:
			self.scheduler.step()

			if accelerator.is_local_main_process:
				wandb.log({'train_loss': m_l1_loss.avg, 'train_ssim': m_ssim.avg, 'train_psnr': m_psnr.avg})

			accelerator.wait_for_everyone()

			valid_epoch_freq = self.opts['valid_epoch_freq'] if 'valid_epoch_freq' in self.opts else 1
			if epoch==start_epoch+1 or (epoch % valid_epoch_freq == 0):
				val_loss, val_ssim, val_psnr = self.validate(epoch, accelerator, run)
				metrics_dict = {'val_loss':val_loss, 'val_ssim':val_ssim, 'val_psnr':val_psnr}

			checkpoint_dict = {'epoch': epoch, 'model': accelerator.unwrap_model(self.net).state_dict(),
							   'optimizer': self.optimizer.state_dict(), 'lr_scheduler': self.scheduler.state_dict(), 'metrics': metrics_dict}

			if epoch % self.opts['save_epoch_freq'] == 0:
				accelerator.save(checkpoint_dict, os.path.join(self.checkpoint_dir, f'checkpoint_epoch_{epoch}.pth'))

			if epoch == start_epoch or (epoch % valid_epoch_freq == 0) or (end_epoch - epoch < 5):
				update_best = val_ssim > self.best_ssim
				if update_best:
					self.best_ssim = val_ssim
					accelerator.print(f'Best valid ssim {self.best_ssim} saved at epoch {epoch}')
					accelerator.save(checkpoint_dict, os.path.join(self.checkpoint_dir, f'checkpoint_best.pth'))
			# save last
			accelerator.save(checkpoint_dict, os.path.join(self.checkpoint_dir, f'checkpoint_last.pth'))

		if accelerator.is_local_main_process:
			wandb.finish()

	@torch.no_grad()
	def validate(self, epoch, accelerator, run):
		self.net.eval()
		batch_time = AverageMeter('Time', ':6.3f', Summary.NONE)
		m_l1_loss = AverageMeter('Loss', ':.4e', Summary.NONE)
		m_ssim = AverageMeter('SSIM', ':6.2f',  Summary.AVERAGE)
		m_psnr = AverageMeter('PSNR', ':6.2f',  Summary.AVERAGE)

		end = time.time()
		for idx, val_loader in enumerate(self.val_loaders):
			correct_count = 0
			# accelerator.print(f'Validation on dataset {datasets_name[idx + 1]}:')
			#
			# with tqdm(total=len(val_loader.dataset), desc=f'Val on {self.datasets_name[idx + 1]}', unit='img',
			with tqdm(total=len(val_loader), desc=f'Val on {self.datasets_name[idx + 1]}', unit='batch',
					  disable=not accelerator.is_local_main_process) as val_pbar:
				for step, batch in enumerate(val_loader):
					image = batch['opt_cloudy']
					sar = batch['sar']
					label = batch['opt_clear']
					if self.use_id:
						image_id = batch['image_id']

					if self.use_id:
						pred = self.net(image, sar, image_id, accelerator)
					elif self.sar_trans:
						pred = self.net(sar)
					else:
						pred = self.net(image, sar)

					# Gathers tensor and potentially drops duplicates in the last batch
					all_pred, all_label = accelerator.gather_for_metrics((pred, label))

					loss_l1 = self.l1_loss(all_pred, all_label)
					# metrics
					ssim = SSIM(all_pred, all_label).item()  # ?
					# ssim = self.ssim(all_pred, all_label).item()
					psnr = PSNR(all_pred, all_label)

					m_l1_loss.update(loss_l1.item() , image.size(0))  # average ?
					m_ssim.update(ssim, image.size(0))
					m_psnr.update(psnr, image.size(0))

					batch_time.update(time.time() - end)
					end = time.time()

					if accelerator.is_local_main_process:
						# figure
						if step==0 or step % self.opts['valid_visual_step_freq'] == 0:
							img_sample = torch.cat([image.data, pred.data, label.data], -1)  # 按宽拼接
							grid = make_grid(img_sample, nrow=1, normalize=True) # 每一行显示的图像列数
							save_image(grid, os.path.join(self.result_dir, 'valid_images', f'img_epoch_{epoch}_step_{step}.png'))

						# val_pbar.update(all_label.shape[0])
						val_pbar.set_postfix(
							ordered_dict={'loss': m_l1_loss.avg, 'ssim': m_ssim.avg, 'psnr': m_psnr.avg})
						val_pbar.update()

		accelerator.wait_for_everyone()
		if accelerator.is_local_main_process:
			wandb.log({'val_loss': m_l1_loss.avg, 'val_ssim': m_ssim.avg, 'val_psnr': m_psnr.avg})

		accelerator.print(f'val_loss: {m_l1_loss.avg}, val_ssim: {m_ssim.avg}, val_ssim: {m_psnr.avg}')
		return m_l1_loss.avg, m_ssim.avg, m_psnr.avg