DeOldify模型精调教程：使用自定义数据集提升上色效果

DeOldify模型精调教程：使用自定义数据集提升上色效果

老照片上色，听起来是个挺酷的事儿。市面上有不少现成的工具，比如DeOldify，效果确实不错。但不知道你有没有遇到过这种情况：想给一批特定风格的老照片上色，比如全是古建筑或者特定年代的人物肖像，用通用模型出来的颜色，总觉得差了那么点意思，要么颜色太艳，要么风格不对味。

这就是通用模型的局限了。它学的是海量通用数据，很难照顾到每个细分领域的独特审美和色彩规律。这时候，如果你手头正好有一批特定领域的黑白-彩色配对照片，自己动手精调一下模型，效果往往会好得多。

今天这篇教程，就是带你走一遍这个精调过程。咱们不搞复杂的理论，就聚焦一件事：怎么用你自己的数据集，让DeOldify模型在你关心的领域里，上色效果更准、更自然。整个过程我会在星图平台的GPU环境里演示，步骤清晰，代码都给出来，你跟着做就行。

1. 精调准备：理解我们要做什么

在开始敲代码之前，咱们先花几分钟把思路理清楚。精调（Fine-tuning）不是从头训练一个模型，那太费时费力了。我们的策略是“站在巨人的肩膀上”。

你可以把预训练好的DeOldify模型想象成一个已经学了大量通用绘画技巧的画师。它知道天空通常是蓝的，树叶是绿的，皮肤是什么色调。现在，我们想请这位画师专门学习某一类画作，比如中国古建筑水彩画。我们不需要它从素描学起，而是给它看大量古建筑的水彩画原作（我们的自定义数据集），让它调整自己的笔触和用色习惯，更好地模仿这类作品的风格。

具体到DeOldify，其核心是一个生成模型。我们精调的主要部分，通常是它的生成器（Generator），这是一个基于U-Net等结构的网络。我们的目标就是通过我们提供的“专业教材”（自定义数据集），微调这个生成器网络的参数，让它输出的颜色更符合我们期望领域的特征。

2. 环境搭建与数据准备

工欲善其事，必先利其器。我们先得把训练环境搭起来，并把“教材”准备好。

2.1 星图GPU环境配置

精调模型是个计算密集型任务，用CPU跑会慢到怀疑人生。星图平台提供了现成的GPU环境，我们直接使用就好。

1. 访问星图镜像广场：在星图平台找到“镜像广场”，搜索“PyTorch”或“深度学习”基础镜像。选择一个包含较新版本PyTorch（如1.12+）、CUDA和常用深度学习库（如torchvision,fastai）的镜像。DeOldify原项目基于fastai，所以包含它的镜像会省事很多。
2. 创建实例：选择该镜像，创建一个GPU实例（例如T4或V100规格）。系统会自动配置好驱动和环境。
3. 连接并验证环境：通过JupyterLab或SSH连接到你的实例。打开一个终端，运行以下命令验证关键组件：

python -c "import torch; print(f'PyTorch版本: {torch.__version__}')" python -c "import torch; print(f'CUDA是否可用: {torch.cuda.is_available()}')" python -c "import fastai; print(f'fastai版本: {fastai.__version__}')"

如果都能正确输出版本信息且CUDA可用，环境就妥了。

2.2 获取DeOldify源码与预训练模型

我们在工作目录下操作：

# 1. 克隆DeOldify仓库（这里以官方仓库为例） git clone https://github.com/jantic/DeOldify.git cd DeOldify # 2. 安装必要的Python依赖 # 官方requirements.txt可能不全，我们补充安装一些 pip install -r requirements.txt pip install jupyterlab wandb # wandb用于可视化，可选但推荐 # 3. 下载预训练模型权重 # DeOldify提供了不同类型的预训练模型，我们选择‘artistic’版本，色彩通常更生动 mkdir -p models wget https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth -O ./models/ColorizeArtistic_gen.pth

2.3 准备自定义数据集

这是精调成败的关键。你需要准备一个黑白-彩色配对的数据集。假设我们专注“古建筑”领域。

数据集结构要求：

your_dataset/ ├── train/ │ ├── blackwhite/ # 存放黑白训练图片 │ │ ├── building_001_bw.jpg │ │ ├── building_002_bw.jpg │ │ └── ... │ └── color/ # 存放对应的彩色训练图片（真值） │ ├── building_001_color.jpg │ ├── building_002_color.jpg │ └── ... └── valid/ ├── blackwhite/ # 存放黑白验证图片 └── color/ # 存放对应的彩色验证图片

数据准备建议：

• 来源：可以从博物馆、历史图库网站寻找高清古建筑照片。用图像处理软件（如Photoshop、GIMP）或代码（如PIL、OpenCV）批量生成对应的黑白版本。
• 配对：确保黑白和彩色图片的文件名一一对应（如building_001_bw.jpg对应building_001_color.jpg），尺寸完全一致。
• 数量：精调不需要像预训练那样百万级的数据。对于古建筑这类特定领域，有几百到几千对质量较高的图片，就能看到明显效果提升。
• 预处理：将图片统一缩放到一个合适的尺寸，例如256x256或512x512。DeOldify模型输入通常为256x256。你可以写一个简单的脚本进行批量处理。

这里提供一个用Python PIL库批量生成黑白图片的示例：

from PIL import Image import os def convert_to_bw(color_dir, bw_dir): """将彩色目录下的图片转换为黑白图，并保存到指定目录""" os.makedirs(bw_dir, exist_ok=True) for img_name in os.listdir(color_dir): if img_name.lower().endswith(('.png', '.jpg', '.jpeg')): color_path = os.path.join(color_dir, img_name) bw_path = os.path.join(bw_dir, img_name.replace('.', '_bw.')) # 修改文件名以区分 try: with Image.open(color_path) as img: bw_img = img.convert('L') # ‘L’模式为灰度图 bw_img.save(bw_path) print(f"Converted: {img_name}") except Exception as e: print(f"Error processing {img_name}: {e}") # 假设你的原始彩色图在 ./raw_color/train 和 ./raw_color/valid convert_to_bw('./raw_color/train', './your_dataset/train/blackwhite') convert_to_bw('./raw_color/valid', './your_dataset/valid/blackwhite') # 注意：你需要手动将原始彩色图复制到 ./your_dataset/train/color 和 ./your_dataset/valid/color

3. 模型精调实战

环境数据都齐了，现在进入核心环节。DeOldify源码中提供了训练脚本，我们需要对其进行一些修改以适应精调任务。

3.1 修改训练配置文件与脚本

首先，我们创建一个自己的精调脚本，比如叫finetune_historical.py。这个脚本基于原项目的train.py，但简化并调整了参数。

# finetune_historical.py import torch from fastai.vision import * from deoldify import visualize from deoldify.loss import * from deoldify.model import * from deoldify.dataset import * import warnings warnings.filterwarnings("ignore") # 1. 基础设置 path = Path('./your_dataset') # 你的数据集路径 bs = 4 # 根据你的GPU内存调整批次大小 size = 256 # 输入图像尺寸 arch = models.resnet34 # 编码器架构，与预训练模型匹配 # 2. 创建数据加载器 (DataBunch) def get_data(bs, size): src = (ImageImageList.from_folder(path/'train', folder_black='blackwhite', folder_color='color') .split_by_folder(train='train', valid='valid') .label_from_func(lambda x: path/'train'/‘color’/x.name.replace(‘_bw’, ‘_color’)) # 根据你的命名调整 .transform(get_transforms(), size=size, tfm_y=True)) return src.databunch(bs=bs).normalize(imagenet_stats, do_y=True) data = get_data(bs, size) # 3. 创建模型并加载预训练权重 learn = unet_learner(data, arch, loss_func=FeatureLoss(), blur=True, norm_type=NormType.Weight) # 加载我们之前下载的预训练生成器权重 learn.model.load_state_dict(torch.load('./models/ColorizeArtistic_gen.pth'), strict=False) print("预训练权重加载完毕。") # 4. 寻找合适的学习率 learn.lr_find() learn.recorder.plot(suggestion=True) # 观察输出图表，选择一个损失仍在明显下降但未陡增的学习率，例如 1e-5 到 1e-4 lr = 1e-4 # 5. 开始精调训练 # 首先微调最后一两层 learn.freeze() learn.fit_one_cycle(5, slice(lr)) print("第一阶段（冻结部分）训练完成。") # 然后解冻所有层，以更小的学习率进行全网络微调 learn.unfreeze() learn.fit_one_cycle(10, slice(lr/10, lr)) print("第二阶段（全网络微调）训练完成。") # 6. 保存精调后的模型 torch.save(learn.model.state_dict(), './models/ColorizeHistorical_gen_finetuned.pth') print("精调模型已保存。")

关键参数解释：

• bs(batch size)：一次训练输入的图片数量。越大训练越快，但需要更多GPU内存。如果内存不足，先调小这个值。
• lr(learning rate)：学习率，控制模型参数更新的步长。精调时通常设置得比从头训练小很多（例如1e-4），以免“冲掉”预训练模型学到的通用知识。
• freeze()&unfreeze()：这是精调的常用技巧。先freeze()大部分底层网络（它们提取通用特征），只训练顶层的少数几层，让模型初步适应新数据。然后unfreeze()所有层，用更小的学习率整体微调。
• fit_one_cycle：一种训练策略，让学习率在一个周期内先增后减，有助于模型更稳定地收敛。

3.2 启动训练与监控

在终端运行你的脚本：

cd /path/to/DeOldify python finetune_historical.py

训练开始后，你需要关注两个东西：

1. 控制台输出：会显示每个epoch的训练损失和验证损失。理想情况下，这两个损失都应该随着训练进行而稳步下降。
2. 验证集效果可视化：最好在训练过程中定期查看模型在验证集上的上色效果。你可以在训练脚本里每隔几个epoch就添加一次可视化代码，或者使用像Weights & Biases (wandb)这样的工具来实时记录损失和生成图片。

一个简单的中间验证示例：

# 在训练循环中适当位置添加，例如每2个epoch if epoch % 2 == 0: learn.show_results(rows=2, figsize=(12,8)) # 或者保存到文件 # save_some_results(learn, path/'valid'/‘blackwhite’)

3.3 应对常见训练问题

• 损失不降或震荡：学习率可能太大了。尝试将lr降低一个数量级（例如从1e-4降到1e-5）。也可能是批次太小，尝试在内存允许下增大bs。
• 过拟合：训练损失持续下降，但验证损失在中后期开始上升。这说明模型只“死记硬背”了训练集。解决方案：增加数据量；使用数据增强（get_transforms()中已包含一些）；早点停止训练（Early Stopping）；在unfreeze()阶段使用更小的学习率。
• 显存不足（OOM）：减小bs和size。也可以尝试使用梯度累积（Gradient Accumulation）来模拟更大的批次。

4. 模型导出与部署

训练完成后，我们得到了一个.pth文件，里面是精调后的模型权重。但怎么用它呢？

4.1 测试精调模型

写一个简单的测试脚本，对比精调前后的效果：

# test_finetuned.py import torch from PIL import Image from deoldify.visualize import * import warnings warnings.filterwarnings("ignore") # 初始化着色器 colorizer = get_image_colorizer(artistic=True) # 加载我们精调后的权重（注意：需要替换colorizer内部的模型权重） # 这里演示一种直接替换的方法，更规范的做法是修改colorizer的初始化逻辑 from deoldify.model import * state_dict = torch.load('./models/ColorizeHistorical_gen_finetuned.pth', map_location=torch.device('cpu')) colorizer.model.load_state_dict(state_dict) colorizer.model.eval() # 测试图片 test_bw_path = './your_dataset/valid/blackwhite/building_test_bw.jpg' # 使用原模型着色（如果需要对比） # result_original = colorizer.get_transformed_image(test_bw_path, render_factor=35) # 使用精调模型着色 colorizer.render_factor = 35 # 渲染因子，影响细节，可调 result_finetuned = colorizer.get_transformed_image(test_bw_path) # 显示或保存结果 result_finetuned.save('./result_finetuned.jpg') print("精调模型上色结果已保存。")

4.2 创建可部署的镜像

为了让别人也能方便地使用你精调好的模型，最好的方式是将其打包成一个Docker镜像。

1. 创建Dockerfile：在DeOldify项目根目录创建一个Dockerfile。

# 基于一个轻量级的PyTorch镜像 FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime # 设置工作目录 WORKDIR /app # 复制项目文件 COPY . . # 安装依赖 RUN pip install --no-cache-dir -r requirements.txt jupyterlab RUN apt-get update && apt-get install -y libgl1-mesa-glx && rm -rf /var/lib/apt/lists/* # 将精调好的模型复制到指定位置（假设在构建镜像前已放在当前目录的models/下） COPY ./models/ColorizeHistorical_gen_finetuned.pth /app/models/ # 暴露端口（例如给Jupyter Lab） EXPOSE 8888 # 设置一个简单的启动命令，例如启动一个推理API服务或Jupyter # 这里示例启动一个简单的HTTP服务（你需要自己实现app.py） # CMD ["python", "app.py"] # 或者直接启动Jupyter Lab CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=''"]

2. 构建并推送镜像：在本地或星图平台的镜像构建服务中，构建这个Docker镜像，并推送到你的镜像仓库。
3. 在星图部署：在星图平台，选择“自定义镜像”部署，填入你的镜像地址，配置好GPU资源和端口，即可一键部署一个包含你精调模型的在线服务。

5. 总结与后续建议

走完这一整套流程，你应该已经得到了一个在古建筑（或你选择的领域）上色效果更专业的DeOldify模型。精调的关键在于数据和耐心：数据要尽可能配对精准、质量高；训练时要细心观察损失曲线，适时调整学习率等参数。

用下来感觉，精调后的模型在特定类型图片上的色彩还原度和风格一致性上，确实比通用模型强不少。当然，这也不是一劳永逸的。如果你的数据集中包含了多种差异很大的子类（比如既有徽派建筑又有故宫），模型可能还需要更精细的设计或更多的数据。

下一步，你可以尝试：

• 探索不同的损失函数：除了DeOldify默认的FeatureLoss，也可以试试结合L1、L2或感知损失，看看对色彩平滑度或细节保留有没有帮助。
• 尝试更大的模型：如果你有更多数据和更强的算力，可以尝试用ResNet50或101作为编码器，模型容量更大，可能学习更复杂的色彩映射。
• 数据增强的学问：针对老照片常见的噪声、划痕、模糊等问题，设计特定的数据增强策略，可以让模型更鲁棒。
• 部署为API服务：将你的精调模型封装成RESTful API，集成到自己的老照片修复应用或网站中，让更多人能用上。

整个过程最有成就感的一刻，莫过于看到一张黑白古建筑照片，被你精调过的模型赋予了恰到好处、符合历史感的色彩。希望这篇教程能帮你走到这一步。如果在实践过程中遇到问题，多看看控制台报错、多查查社区讨论，大部分技术问题都能找到答案。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。