205 lines
6.4 KiB
Python
205 lines
6.4 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
CUT L通道 训练脚本
|
|
==================
|
|
仅对车牌 L (亮度) 通道做纹理迁移训练，a,b 颜色通道不动。
|
|
|
|
用法:
|
|
python train.py \
|
|
--domain-a ./gan_data/domain_A_synthetic/police \
|
|
--domain-b ./gan_data/domain_B_real \
|
|
--epochs 100 \
|
|
--batch 4 \
|
|
--output ./output/cut_police
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
|
|
# 添加 tools 目录到路径
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
|
|
from models import (
|
|
LChannelGenerator, PatchGANDiscriminator,
|
|
PatchNCELoss, init_weights, count_params,
|
|
)
|
|
from dataset import make_loader
|
|
|
|
|
|
def train(args):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
# ---- 数据 ----
|
|
loader_a = make_loader(args.domain_a, args.batch, (args.img_w, args.img_h),
|
|
is_real=False, num_workers=args.workers)
|
|
loader_b = make_loader(args.domain_b, args.batch, (args.img_w, args.img_h),
|
|
is_real=True, num_workers=args.workers)
|
|
|
|
print(f"Domain A (合成): {len(loader_a.dataset)} 张")
|
|
print(f"Domain B (真实): {len(loader_b.dataset)} 张")
|
|
print(f"Batch size: {args.batch}")
|
|
print(f"Epochs: {args.epochs}")
|
|
|
|
# ---- 模型 ----
|
|
G = LChannelGenerator(n_res=args.n_res).to(device)
|
|
D = PatchGANDiscriminator().to(device)
|
|
|
|
init_weights(G)
|
|
init_weights(D)
|
|
|
|
print(f"G params: {count_params(G):,}")
|
|
print(f"D params: {count_params(D):,}")
|
|
|
|
# ---- 损失 & 优化器 ----
|
|
criterion_gan = nn.MSELoss() # LSGAN
|
|
criterion_nce = PatchNCELoss(
|
|
nce_layers=args.nce_layers,
|
|
nce_t=args.nce_t,
|
|
n_patches=args.n_patches,
|
|
).to(device)
|
|
criterion_idt = nn.L1Loss()
|
|
|
|
opt_G = optim.Adam(G.parameters(), lr=args.lr, betas=(0.5, 0.999))
|
|
opt_D = optim.Adam(D.parameters(), lr=args.lr, betas=(0.5, 0.999))
|
|
|
|
# ---- 日志 ----
|
|
os.makedirs(args.output, exist_ok=True)
|
|
|
|
# ---- 训练循环 ----
|
|
for epoch in range(1, args.epochs + 1):
|
|
epoch_start = time.time()
|
|
G_loss_total = 0.0
|
|
D_loss_total = 0.0
|
|
n_batches = 0
|
|
|
|
for real_a, real_b in zip(loader_a, loader_b):
|
|
real_a = real_a.to(device) # (B, 1, H, W) 合成L
|
|
real_b = real_b.to(device) # (B, 1, H, W) 真实L
|
|
B = real_a.size(0)
|
|
|
|
# ---- 1. 更新 D ----
|
|
fake_b = G(real_a).detach()
|
|
# G 不用返回特征
|
|
|
|
pred_real = D(real_b)
|
|
pred_fake = D(fake_b)
|
|
|
|
loss_D_real = criterion_gan(pred_real,
|
|
torch.ones_like(pred_real))
|
|
loss_D_fake = criterion_gan(pred_fake,
|
|
torch.zeros_like(pred_fake))
|
|
loss_D = (loss_D_real + loss_D_fake) * 0.5
|
|
|
|
opt_D.zero_grad()
|
|
loss_D.backward()
|
|
opt_D.step()
|
|
|
|
# ---- 2. 更新 G ----
|
|
fake_b, feat_q = G(real_a, return_features=True)
|
|
|
|
# GAN loss
|
|
pred_fake = D(fake_b)
|
|
loss_G_gan = criterion_gan(pred_fake,
|
|
torch.ones_like(pred_fake))
|
|
|
|
# PatchNCE loss (对比学习)
|
|
_, feat_k = G(real_b, return_features=True)
|
|
loss_G_nce = criterion_nce(feat_q, feat_k)
|
|
|
|
# Identity loss: 裁剪输出到输入尺寸（UNet可能多出几像素）
|
|
identity = G(real_a)
|
|
if identity.shape != real_a.shape:
|
|
identity = identity[:, :, :real_a.shape[2], :real_a.shape[3]]
|
|
loss_G_idt = criterion_idt(identity, real_a) * args.lambda_idt
|
|
|
|
loss_G = loss_G_gan + loss_G_nce * args.lambda_nce + loss_G_idt
|
|
|
|
opt_G.zero_grad()
|
|
loss_G.backward()
|
|
opt_G.step()
|
|
|
|
G_loss_total += loss_G.item()
|
|
D_loss_total += loss_D.item()
|
|
n_batches += 1
|
|
|
|
# ---- Epoch 统计 ----
|
|
G_avg = G_loss_total / max(n_batches, 1)
|
|
D_avg = D_loss_total / max(n_batches, 1)
|
|
elapsed = time.time() - epoch_start
|
|
|
|
print(f"Epoch {epoch:3d}/{args.epochs} | "
|
|
f"G: {G_avg:.4f} D: {D_avg:.4f} | "
|
|
f"{elapsed:.1f}s")
|
|
|
|
pass # 日志仅print，不使用tensorboard
|
|
|
|
# ---- 保存 ----
|
|
if epoch % args.save_every == 0:
|
|
ckpt = {
|
|
'epoch': epoch,
|
|
'G_state': G.state_dict(),
|
|
'D_state': D.state_dict(),
|
|
'opt_G': opt_G.state_dict(),
|
|
'opt_D': opt_D.state_dict(),
|
|
}
|
|
torch.save(ckpt, os.path.join(args.output, f'epoch_{epoch}.pth'))
|
|
|
|
# 最终保存
|
|
final_ckpt = {
|
|
'G_state': G.state_dict(),
|
|
'D_state': D.state_dict(),
|
|
}
|
|
torch.save(final_ckpt, os.path.join(args.output, 'final.pth'))
|
|
print(f"\n训练完成! 模型保存在: {args.output}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='CUT L通道训练')
|
|
|
|
# 数据
|
|
parser.add_argument('--domain-a', type=str, required=True,
|
|
help='合成车牌目录')
|
|
parser.add_argument('--domain-b', type=str, required=True,
|
|
help='真实车牌目录')
|
|
|
|
# 训练
|
|
parser.add_argument('--epochs', type=int, default=100)
|
|
parser.add_argument('--batch', type=int, default=4)
|
|
parser.add_argument('--lr', type=float, default=2e-4)
|
|
parser.add_argument('--workers', type=int, default=2)
|
|
|
|
# 模型
|
|
parser.add_argument('--n-res', type=int, default=6,
|
|
help='残差块数量')
|
|
|
|
# 损失权重
|
|
parser.add_argument('--lambda-nce', type=float, default=1.0)
|
|
parser.add_argument('--lambda-idt', type=float, default=0.5)
|
|
|
|
# PatchNCE
|
|
parser.add_argument('--nce-layers', type=int, nargs='+', default=[0, 1, 2])
|
|
parser.add_argument('--nce-t', type=float, default=0.07)
|
|
parser.add_argument('--n-patches', type=int, default=128)
|
|
|
|
# 保存
|
|
parser.add_argument('--output', type=str, default='./output/cut_train')
|
|
parser.add_argument('--img-w', type=int, default=220)
|
|
parser.add_argument('--img-h', type=int, default=70)
|
|
parser.add_argument('--save-every', type=int, default=20)
|
|
|
|
args = parser.parse_args()
|
|
train(args)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|