DDColor模型微调：适应特定风格的迁移学习技巧

DDColor模型微调：适应特定风格的迁移学习技巧

你是不是遇到过这样的情况：用DDColor给黑白照片上色，效果确实不错，但总觉得色彩风格太“标准”了？比如你想给一组复古漫画上色，希望色彩更鲜艳、更有漫画感，或者想给一批老建筑照片上色，希望色彩更沉稳、更有历史感。这时候，通用的预训练模型就显得有点力不从心了。

今天咱们就来聊聊怎么对DDColor进行微调，让它能适应你想要的特定风格。这就像给一个通用厨师做特训，让他专门学会做你家乡菜的味道。整个过程其实没有想象中那么复杂，跟着步骤走，你也能训练出属于自己的风格化上色模型。

1. 微调前准备：理解DDColor的工作原理

在开始动手之前，先简单了解一下DDColor是怎么工作的，这样后面调整起来心里更有底。

DDColor的核心思路其实挺巧妙的。它不像传统方法那样直接预测每个像素的颜色，而是先学习一组“颜色令牌”——你可以把它们想象成调色盘上的基础颜色。模型会根据输入的黑白图像，从这些颜色令牌中挑选合适的颜色，然后组合起来给图像上色。

这个设计有个很大的好处：微调时我们不需要从头训练整个模型，只需要调整那组颜色令牌的学习方式，让它们更偏向我们想要的风格就行。这就大大降低了训练难度和所需的数据量。

2. 构建你的风格数据集

数据集是微调成功的关键。这里没有统一的标准，完全取决于你想让模型学会什么风格。

2.1 确定风格方向

先想清楚你想要什么效果。比如：

• 动漫风格：色彩鲜艳、对比度高、阴影区域色彩饱和
• 复古胶片风格：色彩偏暖、略带颗粒感、高光柔和
• 水彩画风格：色彩通透、边缘有晕染效果
• 特定艺术家风格：比如莫奈的印象派色彩、梵高的鲜艳笔触

确定了方向，收集数据时就有了明确目标。

2.2 数据收集与处理

理想情况下，你需要成对的图像：一张黑白版，一张彩色版（体现你想要的风格）。但现实中很难找到完全匹配的成对数据，这时候可以灵活处理。

方法一：自己制作训练对如果你有现成的彩色图像，可以用代码批量转换成黑白图：

import cv2 import os from PIL import Image import numpy as np def create_training_pairs(color_dir, output_dir): """ 将彩色图像转换为黑白-彩色对 """ os.makedirs(os.path.join(output_dir, 'bw'), exist_ok=True) os.makedirs(os.path.join(output_dir, 'color'), exist_ok=True) for filename in os.listdir(color_dir): if filename.lower().endswith(('.jpg', '.jpeg', '.png')): # 读取彩色图像 color_path = os.path.join(color_dir, filename) color_img = cv2.imread(color_path) if color_img is None: continue # 转换为黑白 bw_img = cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY) bw_img = cv2.cvtColor(bw_img, cv2.COLOR_GRAY2BGR) # 转回3通道 # 保存 base_name = os.path.splitext(filename)[0] cv2.imwrite(os.path.join(output_dir, 'bw', f'{base_name}.png'), bw_img) cv2.imwrite(os.path.join(output_dir, 'color', f'{base_name}.png'), color_img) print(f'处理完成: {filename}') # 使用示例 create_training_pairs('my_color_images', 'training_data')

方法二：使用现有数据集调整如果找不到足够的数据，可以考虑用现有数据集，但通过后期处理赋予特定风格。比如用ImageNet数据集，但训练时加入风格损失函数。

2.3 数据量要求

对于风格微调，数据量不需要特别大：

• 基础风格：100-200张高质量图像对通常就够了
• 复杂风格：300-500张能获得更好效果
• 关键：质量比数量更重要，确保每张图像都能代表你想要风格

2.4 数据预处理

DDColor训练时需要特定的数据格式。准备好图像后，需要生成元数据文件：

import os import json def generate_meta_file(data_dir, output_file='meta_info.txt'): """ 生成DDColor训练所需的元数据文件 """ bw_dir = os.path.join(data_dir, 'bw') color_dir = os.path.join(data_dir, 'color') meta_lines = [] # 获取所有图像对 bw_files = sorted([f for f in os.listdir(bw_dir) if f.lower().endswith(('.png', '.jpg'))]) for bw_file in bw_files: color_file = bw_file # 假设文件名相同 bw_path = os.path.join('bw', bw_file) color_path = os.path.join('color', color_file) # 检查彩色图像是否存在 if os.path.exists(os.path.join(data_dir, color_path)): # 格式: 彩色图像路径 黑白图像路径 meta_lines.append(f'{color_path} {bw_path}') # 保存到文件 with open(os.path.join(data_dir, output_file), 'w') as f: f.write('\n'.join(meta_lines)) print(f'生成元数据文件完成，共 {len(meta_lines)} 对图像') return len(meta_lines) # 使用示例 image_count = generate_meta_file('training_data') print(f'可用于训练的图像对数量: {image_count}')

3. 微调环境搭建

现在来搭建训练环境。如果你之前用过DDColor，很多步骤应该很熟悉。

3.1 基础环境配置

# 创建conda环境（如果还没有） conda create -n ddcolor_finetune python=3.9 conda activate ddcolor_finetune # 安装PyTorch（根据你的CUDA版本选择） pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 # 安装DDColor依赖 git clone https://github.com/piddnad/DDColor.git cd DDColor pip install -r requirements.txt # 安装训练额外依赖 pip install -r requirements.train.txt # 安装BasicSR python setup.py develop

3.2 下载预训练权重

微调需要从预训练模型开始，这样训练更快、效果更稳定：

import os from modelscope.hub.snapshot_download import snapshot_download def download_pretrained_models(): """ 下载DDColor预训练模型 """ # 创建保存目录 os.makedirs('pretrain', exist_ok=True) print('正在下载预训练模型...') # 下载模型（这里以ddcolor_modelscope为例） model_dir = snapshot_download( 'damo/cv_ddcolor_image-colorization', cache_dir='./pretrain' ) print(f'模型下载完成，保存在: {model_dir}') # 还需要下载ConvNeXt和InceptionV3的预训练权重 # 这些可以从官方链接下载，或者使用以下命令 import requests # ConvNeXt权重 convnext_url = "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth" convnext_path = "./pretrain/convnext_large_22k_224.pth" if not os.path.exists(convnext_path): print('下载ConvNeXt权重...') response = requests.get(convnext_url) with open(convnext_path, 'wb') as f: f.write(response.content) # InceptionV3权重 inception_url = "https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth" inception_path = "./pretrain/inception_v3_google-1a9a5a14.pth" if not os.path.exists(inception_path): print('下载InceptionV3权重...') response = requests.get(inception_url) with open(inception_path, 'wb') as f: f.write(response.content) print('所有预训练权重下载完成！') # 执行下载 download_pretrained_models()

4. 配置微调参数

这是微调的核心步骤，需要根据你的数据和目标调整训练配置。

4.1 修改训练配置文件

DDColor的训练配置在options/train/train_ddcolor.yml。我们需要针对微调做一些调整：

# 数据配置部分 datasets: train: name: MyStyleDataset # 给你的数据集起个名字 type: DDColorDataset dataroot_gt: ./training_data # 你的数据目录 dataroot_lq: ./training_data meta_info_file: ./training_data/meta_info.txt # 上一步生成的元数据文件 io_backend: type: disk # 图像处理参数 gt_size: 256 # 训练图像大小，根据你的数据调整 use_hflip: true # 是否使用水平翻转增强 use_rot: false # 微调时通常不需要旋转 # 批量大小（根据GPU显存调整） batch_size_per_gpu: 8 num_worker_per_gpu: 4 # 网络配置部分 network_g: type: DDColor encoder: convnext decoder: transformer num_queries: 256 # 颜色查询数量，保持默认 num_scales: 3 # 多尺度数量，保持默认 # 预训练权重路径 pretrained: ./pretrain/damo/cv_ddcolor_image-colorization/pytorch_model.pt # 训练配置部分 train: total_iter: 50000 # 总迭代次数，微调可以少一些 warmup_iter: 1000 # 热身迭代 # 优化器配置 optim_g: type: AdamW lr: 1e-5 # 微调时学习率要调小 weight_decay: 0.01 betas: [0.9, 0.999] # 学习率调度 scheduler: type: CosineAnnealingRestartLR periods: [50000] restart_weights: [1] eta_min: 1e-6 # 损失函数权重（微调时可以调整） pixel_weight: 1.0 perceptual_weight: 0.1 style_weight: 0.5 # 增加风格损失权重 adversarial_weight: 0.0 # 微调时通常不用对抗损失 # 验证和保存配置 val: freq: 1000 # 每1000次迭代验证一次 save_checkpoint_freq: 5000 # 每5000次迭代保存一次

4.2 添加风格损失函数（可选但推荐）

如果你想让模型更好地学习特定风格，可以添加风格损失函数。在DDColor代码中添加：

# 在适当的文件中添加以下代码，比如在models/ddcolor_model.py中 import torch import torch.nn as nn import torch.nn.functional as F from torchvision.models import vgg19 class StyleLoss(nn.Module): """ 风格损失函数，基于VGG特征图的Gram矩阵 """ def __init__(self): super().__init__() # 加载VGG19 vgg = vgg19(pretrained=True).features # 提取特定层的特征 self.layer_names = ['relu1_1', 'relu2_1', 'relu3_1', 'relu4_1', 'relu5_1'] self.layers = nn.ModuleList() i = 0 for layer in vgg.children(): if isinstance(layer, nn.Conv2d): i += 1 name = f'conv{i}_{layer.kernel_size[0]}' elif isinstance(layer, nn.ReLU): name = f'relu{i}_{layer.inplace}' layer = nn.ReLU(inplace=False) # 替换为inplace=False的版本 elif isinstance(layer, nn.MaxPool2d): name = f'pool{i}' self.layers.append(layer) if name in self.layer_names: break # 冻结VGG参数 for param in self.parameters(): param.requires_grad = False def compute_gram_matrix(self, x): """ 计算Gram矩阵 """ b, c, h, w = x.size() features = x.view(b, c, h * w) gram = torch.bmm(features, features.transpose(1, 2)) return gram / (c * h * w) def forward(self, generated, target): """ 计算风格损失 generated: 生成图像 target: 目标风格图像 """ style_loss = 0.0 # 提取特征 gen_features = generated target_features = target for layer in self.layers: gen_features = layer(gen_features) target_features = layer(target_features) # 计算Gram矩阵 gen_gram = self.compute_gram_matrix(gen_features) target_gram = self.compute_gram_matrix(target_features) # 计算MSE损失 style_loss += F.mse_loss(gen_gram, target_gram) return style_loss # 然后在训练循环中使用 def calculate_style_loss(generated_images, style_images): """ 计算批次图像的风格损失 """ style_loss_fn = StyleLoss() # 确保图像在0-1范围内 generated = torch.clamp(generated_images, 0, 1) style = torch.clamp(style_images, 0, 1) # 计算损失 loss = style_loss_fn(generated, style) return loss

5. 开始微调训练

一切准备就绪，现在可以开始训练了。

5.1 启动训练脚本

DDColor提供了训练脚本，我们稍微修改一下以适应微调：

#!/bin/bash # scripts/train_finetune.sh export CUDA_VISIBLE_DEVICES=0 # 使用第一块GPU python basicsr/train.py \ -opt options/train/train_ddcolor.yml \ --launcher pytorch \ --auto_resume \ --debug

如果一切正常，你会看到类似这样的输出：

[INFO] Start training from iteration 0 [INFO] Iter 0, lr: 1.00e-05, time: 0.512s, data_time: 0.102s [INFO] Iter 0, loss_pix: 0.1523, loss_percep: 0.0452, loss_style: 0.0231 [INFO] Saving models and training states...

5.2 监控训练过程

训练过程中要关注几个关键指标：

1. 损失值下降趋势：正常情况下，损失值应该逐渐下降并趋于平稳
2. 验证集效果：定期在验证集上测试，确保没有过拟合
3. 生成样本质量：每训练一段时间，保存一些生成样本看看效果

你可以用TensorBoard来可视化训练过程：

# 启动TensorBoard tensorboard --logdir experiments/DDColor_MyStyleDataset/tensorboard # 然后在浏览器中访问 http://localhost:6006

5.3 训练技巧和注意事项

学习率策略：

• 微调时学习率要小，通常用1e-5到1e-4
• 如果训练初期损失震荡太大，可以再调小学习率
• 可以使用学习率warmup，让模型先适应新数据

早停策略：

# 简单的早停实现 class EarlyStopping: def __init__(self, patience=10, min_delta=0): self.patience = patience self.min_delta = min_delta self.counter = 0 self.best_loss = None self.early_stop = False def __call__(self, val_loss): if self.best_loss is None: self.best_loss = val_loss elif val_loss > self.best_loss - self.min_delta: self.counter += 1 if self.counter >= self.patience: self.early_stop = True else: self.best_loss = val_loss self.counter = 0 return self.early_stop # 在训练循环中使用 early_stopper = EarlyStopping(patience=20, min_delta=0.001) for epoch in range(max_epochs): # ... 训练代码 ... val_loss = validate(model, val_loader) if early_stopper(val_loss): print("早停触发，停止训练") break

数据增强：

• 微调时数据增强要适度，避免破坏风格特征
• 推荐使用：水平翻转、小幅度的亮度对比度调整
• 避免使用：大幅度的旋转、裁剪、颜色抖动

6. 评估和应用微调后的模型

训练完成后，需要评估模型效果，然后应用到实际场景中。

6.1 模型评估

import cv2 import torch from torchvision import transforms from PIL import Image def evaluate_finetuned_model(model_path, test_images_dir): """ 评估微调后的模型 """ # 加载模型 from ddcolor import DDColor model = DDColor() checkpoint = torch.load(model_path, map_location='cpu') if 'params' in checkpoint: model.load_state_dict(checkpoint['params']) elif 'params_ema' in checkpoint: model.load_state_dict(checkpoint['params_ema']) else: model.load_state_dict(checkpoint) model.eval() model.cuda() # 预处理 transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) # 测试所有图像 results = [] for img_name in os.listdir(test_images_dir): if img_name.lower().endswith(('.jpg', '.png', '.jpeg')): # 读取图像 img_path = os.path.join(test_images_dir, img_name) img = Image.open(img_path).convert('RGB') # 转换为黑白（用于输入） bw_img = img.convert('L').convert('RGB') # 预处理 input_tensor = transform(bw_img).unsqueeze(0).cuda() # 推理 with torch.no_grad(): output = model(input_tensor) # 后处理 output_img = (output[0].cpu().numpy().transpose(1, 2, 0) * 0.5 + 0.5) * 255 output_img = output_img.astype('uint8') # 保存结果 output_path = os.path.join('results', f'colorized_{img_name}') cv2.imwrite(output_path, cv2.cvtColor(output_img, cv2.COLOR_RGB2BGR)) results.append({ 'original': img_name, 'colorized': output_path }) print(f'评估完成，共处理 {len(results)} 张图像') return results # 使用示例 results = evaluate_finetuned_model( 'experiments/DDColor_MyStyleDataset/models/net_g_latest.pth', 'test_images' )

6.2 效果对比分析

评估时可以从几个角度对比：

1. 色彩准确性：与原始彩色图像对比（如果有的话）
2. 风格一致性：生成的色彩是否符合目标风格
3. 细节保留：边缘、纹理等细节是否清晰
4. 自然度：色彩过渡是否自然，有无明显伪影

6.3 实际应用

训练好的模型可以像原始DDColor一样使用：

# 使用微调后的模型进行推理 def colorize_with_finetuned_model(image_path, model_path): """ 使用微调模型为单张图像上色 """ # 加载模型 model = DDColor() checkpoint = torch.load(model_path, map_location='cpu') model.load_state_dict(checkpoint['params_ema']) model.eval() model.cuda() # 读取并预处理图像 img = Image.open(image_path).convert('RGB') bw_img = img.convert('L').convert('RGB') # 调整大小（保持长宽比） max_size = 512 w, h = bw_img.size if max(w, h) > max_size: scale = max_size / max(w, h) new_w, new_h = int(w * scale), int(h * scale) bw_img = bw_img.resize((new_w, new_h), Image.Resampling.LANCZOS) # 转换为tensor transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) input_tensor = transform(bw_img).unsqueeze(0).cuda() # 推理 with torch.no_grad(): output = model(input_tensor) # 转换为图像 output_img = (output[0].cpu().numpy().transpose(1, 2, 0) * 0.5 + 0.5) * 255 output_img = output_img.astype('uint8') # 调整回原始尺寸 output_pil = Image.fromarray(output_img) output_pil = output_pil.resize((w, h), Image.Resampling.LANCZOS) return output_pil # 批量处理 def batch_colorize(input_dir, output_dir, model_path): """ 批量处理目录中的所有图像 """ os.makedirs(output_dir, exist_ok=True) image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))] for i, filename in enumerate(image_files): print(f'处理中 ({i+1}/{len(image_files)}): {filename}') input_path = os.path.join(input_dir, filename) output_path = os.path.join(output_dir, f'colorized_{filename}') try: result = colorize_with_finetuned_model(input_path, model_path) result.save(output_path) except Exception as e: print(f'处理失败 {filename}: {e}') print('批量处理完成！')

7. 常见问题与解决方案

微调过程中可能会遇到一些问题，这里总结一些常见情况和解决方法。

7.1 训练不收敛或效果差

可能原因：

1. 学习率太大
2. 数据量太少或质量太差
3. 数据预处理有问题
4. 损失函数权重设置不合理

解决方案：

# 调试学习率 def find_optimal_lr(model, train_loader, init_lr=1e-7, end_lr=1e-2, num_iter=100): """ 使用学习率范围测试找到合适的学习率 """ optimizer = torch.optim.AdamW(model.parameters(), lr=init_lr) lr_lambda = lambda x: math.exp(x * math.log(end_lr / init_lr) / num_iter) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) losses = [] lrs = [] model.train() for i, batch in enumerate(train_loader): if i >= num_iter: break # 前向传播和损失计算 # ... loss.backward() optimizer.step() scheduler.step() losses.append(loss.item()) lrs.append(optimizer.param_groups[0]['lr']) # 绘制学习率-损失曲线 # 选择损失下降最快的区域作为学习率 return lrs, losses

7.2 过拟合问题

现象：训练集损失很低，但验证集损失很高或效果差

解决方法：

1. 增加数据增强
2. 使用更小的模型或添加Dropout
3. 早停策略
4. 权重衰减

# 添加Dropout到DDColor模型 class DDColorWithDropout(DDColor): def __init__(self, dropout_rate=0.1, **kwargs): super().__init__(**kwargs) # 在适当的位置添加Dropout self.dropout = nn.Dropout2d(dropout_rate) def forward(self, x): # 原始forward逻辑 # 在特征提取后添加dropout features = self.encoder(x) features = self.dropout(features) # 后续处理... return output

7.3 风格迁移不彻底

现象：生成图像只有部分区域符合目标风格

解决方法：

1. 增加风格损失权重
2. 使用多尺度风格损失
3. 在更多层计算风格损失

# 多尺度风格损失 class MultiScaleStyleLoss(nn.Module): def __init__(self, scales=[1.0, 0.5, 0.25]): super().__init__() self.scales = scales self.style_loss = StyleLoss() def forward(self, generated, target): total_loss = 0 for scale in self.scales: # 缩放图像 if scale != 1.0: h, w = generated.shape[2:] new_h, new_w = int(h * scale), int(w * scale) gen_scaled = F.interpolate(generated, size=(new_h, new_w), mode='bilinear') target_scaled = F.interpolate(target, size=(new_h, new_w), mode='bilinear') else: gen_scaled = generated target_scaled = target # 计算损失 loss = self.style_loss(gen_scaled, target_scaled) total_loss += loss return total_loss / len(self.scales)

8. 进阶技巧与优化

如果你已经掌握了基础微调，可以尝试这些进阶技巧来进一步提升效果。

8.1 渐进式微调

先在大数据集上微调，再在小数据集上精调：

def progressive_finetuning(model, datasets, lr_schedule): """ 渐进式微调 datasets: 列表，每个元素是(数据集, 迭代次数) lr_schedule: 学习率调度计划 """ for dataset_name, dataset, iterations in datasets: print(f'开始微调阶段: {dataset_name}') # 创建数据加载器 train_loader = create_dataloader(dataset) # 设置学习率 current_lr = lr_schedule[dataset_name] optimizer = torch.optim.AdamW(model.parameters(), lr=current_lr) # 训练指定迭代次数 train_for_iterations(model, train_loader, optimizer, iterations) print(f'阶段 {dataset_name} 完成') return model

8.2 注意力机制增强

通过注意力机制让模型更关注风格关键区域：

class SpatialAttention(nn.Module): """空间注意力模块""" def __init__(self, in_channels): super().__init__() self.conv = nn.Conv2d(in_channels, 1, kernel_size=1) self.sigmoid = nn.Sigmoid() def forward(self, x): # 计算注意力权重 attention = self.conv(x) attention = self.sigmoid(attention) # 应用注意力 return x * attention # 在DDColor中集成注意力 class DDColorWithAttention(DDColor): def __init__(self, **kwargs): super().__init__(**kwargs) # 在颜色查询生成前添加注意力 self.spatial_attention = SpatialAttention(256) def forward(self, x): # 原始特征提取 features = self.encoder(x) # 应用空间注意力 attended_features = self.spatial_attention(features) # 后续处理使用注意力增强的特征 # ...

8.3 混合风格训练

如果你想融合多种风格，可以尝试混合训练：

class MixedStyleDataset(torch.utils.data.Dataset): """混合风格数据集""" def __init__(self, style_datasets, mix_weights=None): self.style_datasets = style_datasets self.mix_weights = mix_weights or [1.0] * len(style_datasets) # 计算总长度 self.lengths = [len(ds) for ds in style_datasets] self.total_length = sum(self.lengths) # 创建索引映射 self.index_map = [] for i, length in enumerate(self.lengths): self.index_map.extend([(i, j) for j in range(length)]) def __len__(self): return self.total_length def __getitem__(self, idx): dataset_idx, sample_idx = self.index_map[idx] # 获取样本 sample = self.style_datasets[dataset_idx][sample_idx] # 添加风格标签 sample['style_label'] = dataset_idx sample['style_weight'] = self.mix_weights[dataset_idx] return sample # 在训练时根据风格标签调整损失权重 def mixed_style_loss(output, target, style_labels, style_weights): """ 混合风格损失计算 """ total_loss = 0 for i in range(len(style_labels)): # 获取当前样本的风格 style_idx = style_labels[i] weight = style_weights[style_idx] # 计算该风格对应的损失 style_loss = calculate_specific_style_loss( output[i:i+1], target[i:i+1], style_idx ) total_loss += style_loss * weight return total_loss / len(style_labels)

9. 总结

微调DDColor模型以适应特定风格，整个过程就像教一个已经会画画的学生掌握新的绘画风格。关键是要有明确的目标风格、高质量的训练数据，以及合理的训练策略。

从实际经验来看，成功的微调通常有几个共同点：数据质量高且风格一致、学习率设置合理、训练时间足够但不过度。有时候简单的方法反而效果更好，不必一开始就追求复杂的网络结构或损失函数。

训练完成后，别忘了在实际场景中测试你的模型。看看它在你的目标应用上表现如何，有没有达到预期效果。如果效果不理想，可以回到数据收集阶段，补充一些关键样本，或者调整训练参数再试一次。

微调的魅力在于它的灵活性。同一个基础模型，通过不同的微调策略，可以衍生出无数个 specialized 版本，满足各种特定需求。这比从头训练一个新模型要高效得多，也更容易获得好效果。

获取更多AI镜像

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