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_okTrue) os.makedirs(os.path.join(output_dir, color), exist_okTrue) 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_filemeta_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 python3.9 conda activate ddcolor_finetune # 安装PyTorch根据你的CUDA版本选择 pip install torch2.2.0 torchvision0.17.0 torchaudio2.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 develop3.2 下载预训练权重微调需要从预训练模型开始这样训练更快、效果更稳定import os from modelscope.hub.snapshot_download import snapshot_download def download_pretrained_models(): 下载DDColor预训练模型 # 创建保存目录 os.makedirs(pretrain, exist_okTrue) 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(pretrainedTrue).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 fconv{i}_{layer.kernel_size[0]} elif isinstance(layer, nn.ReLU): name frelu{i}_{layer.inplace} layer nn.ReLU(inplaceFalse) # 替换为inplaceFalse的版本 elif isinstance(layer, nn.MaxPool2d): name fpool{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 loss5. 开始微调训练一切准备就绪现在可以开始训练了。5.1 启动训练脚本DDColor提供了训练脚本我们稍微修改一下以适应微调#!/bin/bash # scripts/train_finetune.sh export CUDA_VISIBLE_DEVICES0 # 使用第一块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 监控训练过程训练过程中要关注几个关键指标损失值下降趋势正常情况下损失值应该逐渐下降并趋于平稳验证集效果定期在验证集上测试确保没有过拟合生成样本质量每训练一段时间保存一些生成样本看看效果你可以用TensorBoard来可视化训练过程# 启动TensorBoard tensorboard --logdir experiments/DDColor_MyStyleDataset/tensorboard # 然后在浏览器中访问 http://localhost:60065.3 训练技巧和注意事项学习率策略微调时学习率要小通常用1e-5到1e-4如果训练初期损失震荡太大可以再调小学习率可以使用学习率warmup让模型先适应新数据早停策略# 简单的早停实现 class EarlyStopping: def __init__(self, patience10, min_delta0): 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(patience20, min_delta0.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_locationcpu) 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, fcolorized_{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 效果对比分析评估时可以从几个角度对比色彩准确性与原始彩色图像对比如果有的话风格一致性生成的色彩是否符合目标风格细节保留边缘、纹理等细节是否清晰自然度色彩过渡是否自然有无明显伪影6.3 实际应用训练好的模型可以像原始DDColor一样使用# 使用微调后的模型进行推理 def colorize_with_finetuned_model(image_path, model_path): 使用微调模型为单张图像上色 # 加载模型 model DDColor() checkpoint torch.load(model_path, map_locationcpu) 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_okTrue) 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处理中 ({i1}/{len(image_files)}): {filename}) input_path os.path.join(input_dir, filename) output_path os.path.join(output_dir, fcolorized_{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 训练不收敛或效果差可能原因学习率太大数据量太少或质量太差数据预处理有问题损失函数权重设置不合理解决方案# 调试学习率 def find_optimal_lr(model, train_loader, init_lr1e-7, end_lr1e-2, num_iter100): 使用学习率范围测试找到合适的学习率 optimizer torch.optim.AdamW(model.parameters(), lrinit_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, losses7.2 过拟合问题现象训练集损失很低但验证集损失很高或效果差解决方法增加数据增强使用更小的模型或添加Dropout早停策略权重衰减# 添加Dropout到DDColor模型 class DDColorWithDropout(DDColor): def __init__(self, dropout_rate0.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 output7.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), modebilinear) target_scaled F.interpolate(target, size(new_h, new_w), modebilinear) 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(), lrcurrent_lr) # 训练指定迭代次数 train_for_iterations(model, train_loader, optimizer, iterations) print(f阶段 {dataset_name} 完成) return model8.2 注意力机制增强通过注意力机制让模型更关注风格关键区域class SpatialAttention(nn.Module): 空间注意力模块 def __init__(self, in_channels): super().__init__() self.conv nn.Conv2d(in_channels, 1, kernel_size1) 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_weightsNone): 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:i1], target[i:i1], style_idx ) total_loss style_loss * weight return total_loss / len(style_labels)9. 总结微调DDColor模型以适应特定风格整个过程就像教一个已经会画画的学生掌握新的绘画风格。关键是要有明确的目标风格、高质量的训练数据以及合理的训练策略。从实际经验来看成功的微调通常有几个共同点数据质量高且风格一致、学习率设置合理、训练时间足够但不过度。有时候简单的方法反而效果更好不必一开始就追求复杂的网络结构或损失函数。训练完成后别忘了在实际场景中测试你的模型。看看它在你的目标应用上表现如何有没有达到预期效果。如果效果不理想可以回到数据收集阶段补充一些关键样本或者调整训练参数再试一次。微调的魅力在于它的灵活性。同一个基础模型通过不同的微调策略可以衍生出无数个 specialized 版本满足各种特定需求。这比从头训练一个新模型要高效得多也更容易获得好效果。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。