Deraining项目配置与训练指南

安装conda环境

conda create -n vmambair python=3.9
conda activate vmambair

安装PyTorch和相关库

pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121

克隆并安装项目

cd kernels/selective_scan
pip install .

接着进入Deraining目录，运行以下命令安装依赖：

cd Deraining
pip install timm fvcore

创建训练配置文件

# 项目名称或配置名称
name: Deraining_Restormer
# 模型类型
model_type: ImageCleanModel
# 缩放比例（通常用于超分辨率任务，这里是1，表示不缩放）
scale: 1
# 训练时使用的GPU数量，设置为0表示CPU模式
num_gpu: 1
# 随机种子，用于保证实验的可复现性
manual_seed: 100
# 数据集类型
type: Dataset_PairedImage
# 高质量（Ground Truth）图像的路径
dataroot_gt: /root/autodl-tmp/VmambaIR-main/Deraining/Dataset/train/B
# 低质量（输入）图像的路径
dataroot_lq: /root/autodl-tmp/VmambaIR-main/Deraining/Dataset/train/A
# 是否开启几何增强（如旋转、翻转等）
geometric_augs: true
# 文件名模板，用于匹配图像文件名
filename_tmpl: '{}'
# 输入输出后端类型，这里是磁盘存储
io_backend: type: disk
# 数据 loader 设置
use_shuffle: true
num_worker_per_gpu: 4
batch_size_per_gpu: 2
# 渐进式训练设置
mini_batch_sizes: [8,5,4,2,1,1]
iters: [92000,64000,48000,36000,36000,24000]
gt_size: 384
gt_sizes: [128,160,192,256,320,384]
# 训练设置
total_iter: 300
warmup_iter: -1
use_grad_clip: true
# 学习率调度器类型
scheduler: type: CosineAnnealingRestartCyclicLR
periods: [92000, 208000]
restart_weights: [1,1]
eta_mins: [0.0003,0.000001]
# 数据增强设置
mixing_augs: mixup: false
mixup_beta: 1.2
use_identity: true
# 优化器设置
optim_g: type: AdamW
lr: !!float 3e-4
weight_decay: !!float 1e-4
betas: [0.9, 0.999]
# 损失函数设置
pixel_opt: type: L1Loss
loss_weight: 1
reduction: mean
# 验证设置
val: window_size: 8
val_freq: !!float 4e3
save_img: false
rgb2bgr: true
use_image: true
max_minibatch: 8
metrics: psnr: type: calculate_psnr
test_y_channel: true
# 日志设置
logger: print_freq: 1000
save_checkpoint_freq: !!float 4e3
use_tb_logger: true
# 排队训练设置
dist_params: backend: nccl
port: 29500

启动训练命令

python basicsr/train.py -opt /root/autodl-tmp/VmambaIR-main/Deraining/Options/Deraining.yml

单次测试

python test.py

批量测试