PMRF项目安装与配置指南-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00774/article/details/147136680

PMRF项目安装与配置指南

PMRF [ICLR 2025] Official implementation of Posterior-Mean Rectified Flow: Towards Minimum MSE Photo-Realistic Image Restoration 项目地址: https://gitcode.com/gh_mirrors/pm/PMRF

1. 项目基础介绍

PMRF（Posterior-Mean Rectified Flow）是一种新颖的图像复原算法，它通过近似最优估计器来最小化均方误差（MSE），在完美的感知质量约束下工作。该算法的目标是恢复出视觉效果逼真的图像。本项目是PMRF算法的开源实现，提供了算法的代码、训练脚本以及测试数据。

主要编程语言：Python

2. 项目使用的关键技术和框架

PyTorch: 用于深度学习模型的构建和训练。
PyTorch Lightning: 简化PyTorch工作流程的库，用于更高效地训练模型。
natten: 用于实现HDiT模型架构的注意力机制。
timm: 提供了许多预训练模型和模型架构。
opencv-python: 用于图像处理。
wandb: 用于实验跟踪和可视化。

3. 项目安装和配置的准备工作

在开始安装之前，请确保您的系统满足以下要求：

Python 3.10
CUDA 11.8 (如果使用natten包)
conda (推荐使用conda环境管理)

详细安装步骤

创建conda环境

conda create -n pmrf python=3.10
conda activate pmrf

安装PyTorch和相关依赖

conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=11.8 -c pytorch -c nvidia

安装PyTorch Lightning

conda install lightning==2.3.3 -c conda-forge

安装其他Python包

pip install opencv-python==4.10.0.84 timm==1.0.8 wandb==0.17.5 lovely-tensors==0.1.16 torch-fidelity==0.3.0 einops==0.8.0 dctorch==0.1.2 torch-ema==0.3 natten==0.17.1+torch230cu118 -f https://shi-labs.com/natten/wheels nvidia-cuda-nvcc-cu11 basicsr==1.4.2 git+https://github.com/toshas/torch-fidelity.git lpips==0.1.4 piq==0.8.0 huggingface_hub==0.24.5

注意：请确保替换natten==0.17.1+torch230cu118为您系统上安装的正确CUDA版本。

修改basicsr中的代码 由于兼容性问题，需要修改basicsr包中的一个文件。打开以下路径的文件：

/path/to/env/pmrf/lib/python3.10/site-packages/basicsr/data/degradations.py

将以下行：

from torchvision.transforms.functional_tensor import rgb_to_grayscale

修改为：

from torchvision.transforms.functional import rgb_to_grayscale

下载预训练模型和测试数据集 根据项目README中的指示，从提供的链接下载预训练模型和测试数据集，并将其放置在相应的目录中。

完成以上步骤后，您就可以开始使用PMRF项目了。按照项目README中的指示运行训练和测试脚本，以验证安装是否成功。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考