神经网络复现步骤

一、神经网络

1、根据architecture图，记录通道数变化

引入import torch.nn as nn

创建通用卷积块方法

2、定义类，继承(nn.Module):

de __init__中初始化每个卷积块或设定的结构块

注意：此处需要注意级联后图像通道数的变化

def forward(定义参数)

参数：如果是单通道输入，则可设定为x；如是双通道输入，则为两个图，设定为x_1，x_2

forward方法中，根据结构图，源图像依次传入卷积块或设定的结构块中

注意：此处需要注意级联、叠加操作

二、训练

1、创建参数类，写入超参数

可以先创建参数类，将超参数写入参数类中

或使用argparse.ArgumentParser()，调用add_argument方法写入参数

必须参数：数据目录（包括训练集数据、测试集数据目录）、裁剪尺寸(crop/image_size)、batch_size、epoch、learning_rate、训练模型保存目录、训练模型名称、device是cuda还是cpu

可选参数：图像维度dim、stride、log_dir

2、损失函数类

定义所需损失函数，如MS_SSIM、MSE等

3、工具类Utils

包括：数据集读取、文件读取、图像加载、图像裁剪处理、RGB转YCrCb、YCrCb转RGB

注意：如果工具类中没有图像加载功能，则需要创建图像数据加载类加载图像，并对图像做相应裁剪、通道转换等处理

4、创建Train类

（1）创建Dataloader加载数据 / 使用加载图像类工具加载训练数据集

（2）创建模型

（3）定义损失函数

（4）定义优化器：传入model.parameters(), lr=0.0001

（5）加载预训练模型

使用torch.load()    model.load_state_dict(加载的模型)

备注：model.load_state_dict()  用于将预训练的参数权重加载到新的模型中

其中strict= 参数：匹配原权重和现在网络中对应层的权重。当strict为True时，预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；当strict为False时，如果新构建的模型在层数上进行了部分微调，则上述代码就会报错，key对不上

（6）开始训练，定义epoch循环——加载训练图像——图像.to(device)——model(传入源图像)，此处调用的是forward方法——optmizer.zero_grad()——损失计算——loss.backward()反向传播——optimizer.step()——保存训练模型——保存损失数据——绘制损失曲线(绘制曲线使用matplotlib中的pyplot，方法为plot.plt(值, '-')，plot.savefig(路径) 为保存图片)

三、测试

• 定义测试类
• 加载测试图像
• 转换图像通道、尺寸、RGB2YCrCb等
• with torch.no_grad()  测试时无需梯度
• 将测试图像加入模型中，融合
• 融合图像后期处理
• 设置融合图像保存路径，保存融合图像