电话销售-网站建设-开场白,官方网站建设流程及费用,wordpress一键分享,株洲市建设局官方网站Pytorch训练RCAN QAT超分模型 版本信息测试步骤准备数据集创建容器生成文件列表创建文件列表的代码执行脚本,生成文件列表训练RCAN模型准备工作修改开源代码编写训练代码执行训练脚本可视化本文以RCAN超分模型为例,演示了QAT的训练过程,步骤如下: 先训练FP32模型再加载FP32训练… Pytorch训练RCAN QAT超分模型 版本信息测试步骤准备数据集创建容器生成文件列表创建文件列表的代码执行脚本,生成文件列表 训练RCAN模型准备工作修改开源代码编写训练代码执行训练脚本 可视化 本文以RCAN超分模型为例,演示了QAT的训练过程,步骤如下:
先训练FP32模型再加载FP32训练的权值,进行QAT训练连续5次loss没有下降则停止训练为了加快演示,当psnr大于33.0时就停止训练采用tensorboard观察Loss曲线版本信息
属性值训练环境搭建步骤GPU型号NVIDIA GeForce RTX 3080 12GB数据集下载链接http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.ziphttp://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.ziphttp://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.ziphttp://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip开源模型结构https://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/model/rcan.pyhttps://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/option.pyhttps://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/model/common.pyhttps://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/template.py测试步骤
准备数据集
wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip
wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip
wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.zip
wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip创建容器
按https://editor.csdn.net/md/?articleId=136176989的步骤构建镜像
docker stop rcan_dev
docker rm rcan_dev
nvidia-docker run -ti -e NVIDIA_VISIBLE_DEVICES=all --privileged \--net=host -p 6006:6006 -v $PWD:/home -w /home \-v /mnt/disk/RCAN/:/RCAN --name rcan_dev cuda_dev_image:v1.0 /bin/bash
conda activate ai_dev生成文件列表
创建文件列表的代码
# generate_datalist.pyimport os
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdmtrain_HR_path = './DIV2K_train_HR'
train_LR_path = './DIV2K_train_LR_bicubic/X2'
valid_HR_path = './DIV2K_valid_HR'
valid_LR_path = './DIV2K_valid_LR_bicubic/X2'train_file = 'datalist_div2k_train.txt'
valid_file = 'datalist_div2k_valid.txt'def get_images(input_path, format='png'):names = [os.path.splitext(fname)[0]for fname in os.listdir(input_path)if fname.endswith(format)]names.sort()return namesdef get_folders(input_path):names = [directory for directory in os.listdir(input_path)if os.path.isdir(os.path.join(input_path, directory))]names.sort()return namesthe_train_file = open(train_file, 'w')
image_names = get_images(train_HR_path)
for image_name in image_names:the_train_file.write('DIV2K_train_LR_bicubic/X2/' + image_name + 'x2.png' + ' ' + 'DIV2K_train_HR/' + image_name + '.png' + '\n')
the_train_file.close()the_valid_file = open(valid_file, 'w')
image_names = get_images(valid_HR_path)
for image_name in image_names: the_valid_file.write('DIV2K_valid_LR_bicubic/X2/' + image_name + 'x2.png' + ' ' + 'DIV2K_valid_HR/' + image_name + '.png' + '\n')
the_valid_file.close()执行脚本,生成文件列表
cd /RCAN/
unzip DIV2K_train_HR.zip
unzip DIV2K_valid_HR.zip
unzip DIV2K_train_LR_bicubic_X2.zip
unzip DIV2K_valid_LR_bicubic_X2.zip
python generate_datalist.py训练RCAN模型
准备工作
# 安装依赖
pip install tensorboard -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install scikit-image -i https://pypi.tuna.tsinghua.edu.cn/simple# 设置环境变量
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python# 下载开源模型源码
cd /RCAN/
mkdir model
curl -L -o model/rcan.py https://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/model/rcan.py
curl -L -o model/option.py https://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/option.py
curl -L -o model/common.py https://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/model/common.py
curl -L -o template.py https://raw.githubusercontent.com/RussellEven/Multi-frame-RCAN/master/code/template.py修改开源代码
model/rcan.pymodel/common.py 编写训练代码
# train.pyimport os
import torch
import torch.nn as nn
import torch.optim as optim
import json
import copy
import time
from torch.utils.data import DataLoader
from tqdm import tqdm
from torch.quantization.quantize_fx import prepare_qat_fx,convert_fx
from torch.ao.quantization import qconfig
from torch.ao.quantization.fake_quantize import *
from torch.ao.quantization.observer import *
from torch.utils import tensorboard
from torch.autograd import Variable
from torch.utils.data import Dataset
from skimage.color import rgb2hsv, hsv2rgb
import imageio
import random
import numpy as npdef _apply(func, x):if isinstance(x, (list, tuple)):return [_apply(func, x_i) for x_i in x]elif isinstance(x, dict):y = {}for key, value in x.items():y[key] = _apply(func, value)return yelse:return func(x)def get_patch(*args, patch_size=96, scale=2, input_large=False):ih, iw = args[0].shape[:2]if not input_large:p = scaletp = p * patch_sizeip = tp // scaleelse:tp = patch_sizeip = patch_sizeix = random.randrange(0
