基于棋盘格上下文模型的图像压缩算法（zqb第四章-RBs）

前言

此篇博客仅供自己观看，记录实验复现中的问题。
本实验在实验室集群跑

1、数据集准备

数据集取自于 ImageNet Validation dataset，该数据集中包含50000 张 JPEG 图像，从中取最大的 16000 张作为训练集，再取剩余图像中最大的2000 张作为测试集。
λ 的值设置为{0.0016, 0.0032, 0.0075, 0.015, 0.03, 0.045}，对应六个不同码率的模型，前三个模型的超参数设置为 N = 192，M = 192，后三个模型为 N = 192, M = 320。

2、调用train.py进行训练

调用脚本进行训练，调用.sh脚本方法

python3 -u train.py -d /home/disk/lilin/datasets/imageNet --epochs 501 -lr 1e-4 --lambda 0.0016 --batch-size 16 --cuda --save > log.out

3、整理训练流程

1. 读取数据集地址，加载数据集
2. 将数据集分为训练集与测试集，然后随机剪裁数据中的图片为patchsize*patchsize大小
3. 加载构建好的网络模型
4. 将设置好的模型与相应设置的参数传入优化器
5. 判断是否有训练一部分的模型，如果有则继续训练，如果没有则重新开始
6. 若重新开始训练模型，则将数据送入优化器，进行正向传播与反向梯度传播，并计算bpp、mse、loss的均值打印（10次迭代打印一次）计算率失真损失，并保存节点信息

4、与基础分支差别

将激活函数改为了RBs

5、RBs模块

import torch
import torch.nn as nn

def conv1x1(in_ch: int, out_ch: int, stride: int = 1) -> nn.Module:
    '''1x1 convolution'''
    return nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=stride)

def conv3x3(in_ch: int, out_ch: int, stride: int = 1) -> nn.Module:
    '''3x3 convolution with padding'''
    return nn.Conv2d(