论文复现：Pytorch实现VDSR

from torch.utils.data.dataloader import DataLoader #from dataset import TrainDataset from dataset import DatasetFromHdf5 from vdsr import Net import torch.nn as nn import torch.optim as optim from torch.optim.lr_scheduler import MultiStepLR import torch from tqdm import tqdm import os, argparse import utils import matplotlib.pyplot as plt import h5py import numpy as np # parameters setting parser = argparse.ArgumentParser(description="PyTorch VDSR") parser.add_argument("--datasets_path", type=str, default=f'datasets/291_aug_train_VDSRx2.h5', help=".h5 file path") parser.add_argument("--weight_save_path", type=str, default=f'weight/Train_VDSR_model_x2.pth', help="weight file save path") parser.add_argument("--batch_size", type=int, default=64, help="Train batch size, Default: 64") parser.add_argument("--num_workers", type=int, default=16, help="Train num workers, Default: 16") parser.add_argument("--momentum", type=float, default=0.9, help="Momentum, Default: 0.9") parser.add_argument("--weight_decay", type=float, default=0.0001, help="Weight decay, Default: 1e-4") parser.add_argument("--init_lr", type=float, default=0.1, help="Learning Rate, Default=0.1") parser.add_argument("--theta", type=float, default=0.01, help="Clipping Gradients theta, Default: 0.01") parser.add_argument("--epochs", type=int, default=80, help="Number of epochs to train for, Default: 80") parser.add_argument("--step", type=int, default=20, help="Every 20 epochs, lr down 10 times, Default: 20") parser.add_argument("--gamma", type=float, default=0.1, help="lr down times, Default: 0.1") # 设置所有可以使用的显卡，共计四块，根据自己设备修改 os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3' device_ids = [0, 1, 2, 3] opt = parser.parse_args() #print(opt) # 加载数据集 #train_dataset = TrainDataset(opt.datasets_path) train_dataset = DatasetFromHdf5(opt.datasets_path) # print(train_dataset.data[0].shape) #print(train_dataset.target.size) train_dataloader = DataLoader(dataset=train_dataset, batch_size=opt.batch_size, num_workers=opt.num_workers, shuffle=True) # 模型训练相关 model = Net() criterion = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=opt.init_lr, momentum=opt.momentum, weight_decay=opt.weight_decay) scheduler = MultiStepLR(optimizer, milestones=[opt.step, opt.step*2, opt.step*3], gamma=opt.gamma) # 设置GPU model = nn.DataParallel(model, device_ids=device_ids) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # 训练 # TODO 论文中实验部分没有对验证环节说明，但是如果训练中累加psnr，最终现存会溢出。 # best_psnr = 0 # each_psnr = [] for epoch in range(opt.epochs): epoch_loss = 0 #sum_psnr = 0 print("Epoch = {}, lr = {}".format(epoch, optimizer.param_groups[0]["lr"])) with tqdm(train_dataloader, desc='standing...') as tepoch: for iteration, (input, target) in enumerate(tepoch): input, target = input.to(device), target.to(device) optimizer.zero_grad() out = model(input) loss = criterion(out, target) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), opt.theta/optimizer.param_groups[0]["lr"]) optimizer.step() epoch_loss += loss.item() tepoch.set_description(f'Epoch [{epoch + 1}/{opt.epochs}]') #vdsr_psnr = utils.calc_psnr(input, out) #sum_psnr += vdsr_psnr if iteration % 100 == 0: print("===> Epoch[{}]({}/{}): Loss: {:.10f}".format(epoch, iteration, len(train_dataloader), loss.item())) # print("===> Epoch[{}]({}/{}): Loss: {:.10f}, PSNR: {:.2f} ".format(epoch, iteration, len(train_dataloader), # loss.item(), vdsr_psnr)) #avg_psnr = sum_psnr / len(train_dataloader) # print(f"Average PSNR: {avg_psnr} dB") # if epoch > 0: # TODO 忽略第一个epoch，因为第一个epoch是从最大往下降的，值偏高 # each_psnr.append(avg_psnr) # if avg_psnr >= best_psnr and epoch > 0: # best_psnr = avg_psnr # 用psnr衡量模型，保存最好的 # torch.save(model, r"weight/best_model_VDSR_x2.pth") print(f"Epoch {epoch}. Training loss: {epoch_loss / len(train_dataloader)}") scheduler.step() # 训练完保存模型 if len(device_ids) > 1: torch.save(model.module.state_dict(), opt.weight_save_path) else: torch.save(model.state_dict(), opt.weight_save_path) torch.save(model, 'weight/VDSR_model_x2_80epoch.pth') # # 画epoch和psnr的关系 # epochs = range(0, opt.epochs) # # # plot画图，设置颜色和图例，legend设置图例样式，右下，字体大小 # plt.plot(epochs, each_psnr, color='red', label='VDSR(trained on T291)') # plt.legend(loc='lower right', prop={'family':'Times New Roman', 'size': 12}) # # # 横纵轴文字内容以及字体样式 # plt.xlabel("epoch",fontproperties='Times New Roman', size = 12) # plt.ylabel("Average test PSNR (dB)",fontproperties='Times New Roman', size=12) # # # 横纵轴刻度以及字体样式 # plt.xlim(0, opt.epochs) # plt.ylim(20, 40) # plt.yticks(fontproperties='Times New Roman', size=12) # plt.xticks(fontproperties='Times New Roman', size=12) # # # 网格样式 # plt.grid(ls = '--') # plt.savefig('fig4-a.png', bbox_inches='tight') # plt.show() # # # epoch_loss = 0 # # # 查看训练集第一个批次，看数据集是否封装正确 # # for batch_idx, (data, target) in enumerate(train_dataloader): # # # data 和 target 是每个批次的数据和标签 # # # 这里只输出第一个批次的数据 # # print("Batch Index:", batch_idx) # # print("Data shape:", data.shape) # 输出数据形状 # # print("Target shape:", target.shape) # 输出标签形状 # # print("First Batch Data:", data[0]) # 输出第一个批次的数据 # # print("First Batch Target:", target[0]) # 输出第一个批次的标签 # # # # optimizer.zero_grad() # # out = model(data) # # loss = criterion(out, target) # # loss.backward() # # optimizer.step() # # epoch_loss += loss.item() # # print(loss) # # print(epoch_loss) # # print(epoch_loss / len(train_dataloader)) # # break # # # if batch_idx == 3: # # break # 只输出第一个批次 # # # print(input.shape) # # print(label.shape)