Zi2ZiModel 字型風格轉換加入對比學習(Contrastive Loss)

2025-03-072025-03-07

Contrastive Learning (CL) 與 VAE、GAN 類似，都是可以作為 Unsupervised Learning 無監督學習方法的學習方式。但用在無標籤數據時，通常我們會將對比學習歸類為自監督學習 Self-Supervised Learning，而不是無監督學習。

在字型風格轉換任務中，使用 Contrastive loss（對比損失）確實有可能有助於提升學習的穩定性和效果，尤其是在以下幾個方面：

1. 特徵空間的對齊：

Contrastive loss 的核心思想是將相似的樣本在特徵空間中拉近，將不相似的樣本推遠。
在字型風格轉換中，我們可以將同一字符的不同風格版本視為相似樣本，將不同字符的風格版本視為不相似樣本。
通過 Contrastive loss，我們可以促使模型學習到一個能夠將同一字符的不同風格版本映射到相近特徵的特徵空間，從而提高風格轉換的一致性。

2. 風格特徵的提取：

Contrastive loss 可以幫助模型更好地提取風格特徵，並將其與內容特徵分離。
通過將同一風格的不同字符版本在特徵空間中拉近，我們可以促使模型學習到能夠捕捉風格共性的特徵表示。
這有助於提高風格轉換的準確性和泛化能力。

3. 訓練的穩定性：

Contrastive loss 相對於傳統的像素級損失函數，對於噪聲和異常值具有更強的魯棒性。
這有助於提高訓練的穩定性，避免模型在訓練過程中出現崩潰或震盪。
另外對比損失可以幫助模型學習到更有區別性的特徵，這可以讓判別器有更好的判別能力，進而穩定GAN的訓練。

具體應用方式：

在字型風格轉換的生成器中，我們可以添加一個額外的特徵提取器，用於提取字符的風格特徵。
然後，我們可以使用 Contrastive loss 來訓練這個特徵提取器，使其能夠將同一字符的不同風格版本映射到相近的特徵。
最後，我們可以將提取到的風格特徵用於生成器的風格轉換過程。

總結：

使用 Contrastive loss 有潛力提升字型風格轉換學習的穩定性和效果。
通過對齊特徵空間、提取風格特徵和提高訓練穩定性，Contrastive loss 可以幫助模型學習到更準確和魯棒的風格轉換模型。

注意事項：

Contrastive loss 的效果取決於樣本對的選擇和損失函數的設計。
在實際應用中，需要仔細調整相關參數，以獲得最佳效果。
Contrastive loss 的計算複雜度可能較高，需要考慮計算資源的限制。

你可以透過對比學習（Contrastive Loss）來提升字型風格轉換的學習穩定性。以下是你的 trainer script 需要修改的幾個關鍵部分：

步驟 1：加入對比學習 Loss

在 Zi2ZiModel 的 __init__ 方法中，新增對比損失函數：

from pytorch_metric_learning import losses

class Zi2ZiModel:
    def __init__(self, ...):
        ...
        self.contrastive_loss = losses.NTXentLoss(temperature=0.07)  # 溫度超參數可調整
        if self.gpu_ids:
            self.contrastive_loss.cuda()

這裡使用 Normalized Temperature-scaled Cross Entropy Loss (NTXentLoss)，它適合無監督對比學習。

步驟 2：修改 `forward()` 方法

在 self.forward() 方法中，確保 特徵嵌入（feature embeddings） 可以用於對比學習：

def forward(self):
    self.fake_B, self.encoded_real_A = self.netG(self.real_A, self.labels)
    self.encoded_fake_B = self.netG(self.fake_B).view(self.fake_B.shape[0], -1)
    
    # 為了對比學習，我們需要一個正樣本 (real_A) 和一個負樣本 (fake_B)
    self.embedding_real_A = self.encoded_real_A.view(self.encoded_real_A.shape[0], -1)
    self.embedding_fake_B = self.encoded_fake_B.view(self.encoded_fake_B.shape[0], -1)

這樣可以確保 encoded_real_A 和 encoded_fake_B 被展平為向量，適合作為對比學習的輸入。

步驟 3：修改 `backward_G()`，加入對比學習 Loss

在 backward_G() 方法中，加入 對比學習損失：

def backward_G(self, no_target_source=False):
    fake_AB = torch.cat([self.real_A, self.fake_B], 1)
    fake_D_logits, fake_category_logits = self.netD(fake_AB)

    const_loss = self.Lconst_penalty * self.mse(self.encoded_real_A, self.encoded_fake_B)
    l1_loss = self.L1_penalty * self.l1_loss(self.fake_B, self.real_B)
    fake_category_loss = self.Lcategory_penalty * self.category_loss(fake_category_logits, self.labels)
    cheat_loss = self.real_binary_loss(fake_D_logits)

    # 計算對比損失
    contrastive_loss = self.contrastive_loss(self.embedding_real_A, self.embedding_fake_B)
    
    self.g_loss = cheat_loss + l1_loss + fake_category_loss + const_loss + contrastive_loss
    self.g_loss.backward()
    
    return const_loss, l1_loss, cheat_loss, contrastive_loss

這裡 contrastive_loss 會確保 real_A 和 fake_B 的特徵保持區別，使得生成器學習更穩定的字型風格轉換。

步驟 4：修改 `optimize_parameters()`

在 optimize_parameters() 方法的返回值中加入 contrastive_loss：

def optimize_parameters(self):
    self.forward()
    
    # 更新 D
    self.set_requires_grad(self.netD, True)
    self.optimizer_D.zero_grad()
    category_loss = self.backward_D()
    self.optimizer_D.step()

    # 更新 G
    self.set_requires_grad(self.netD, False)
    self.optimizer_G.zero_grad()
    const_loss, l1_loss, cheat_loss, contrastive_loss = self.backward_G()
    self.optimizer_G.step()
    
    return const_loss, l1_loss, category_loss, cheat_loss, contrastive_loss

步驟 5：修改訓練迴圈

最後，在 trainer script 中記錄 contrastive_loss：

for bid, batch in enumerate(dataloader):
    model.set_input(batch[0], batch[2], batch[1])
    const_loss, l1_loss, category_loss, cheat_loss, contrastive_loss = model.optimize_parameters()

    if bid % 100 == 0:
        passed = time.time() - start_time
        log_format = "Epoch: [%2d], [%4d/%4d] time: %4.2f, d_loss: %.5f, g_loss: %.5f, " + \
                     "category_loss: %.5f, cheat_loss: %.5f, const_loss: %.5f, l1_loss: %.5f, contrastive_loss: %.5f"
        print(log_format % (epoch, bid, total_batches, passed, model.d_loss.item(), model.g_loss.item(),
                            category_loss, cheat_loss, const_loss, l1_loss, contrastive_loss))

這樣修改後，對比學習的作用

強化字型風格區分能力：讓 real_A 和 fake_B 在特徵空間保持一定距離，避免生成器產生模糊結果。
提升學習穩定性：與 L1 Loss、對抗損失結合，減少模式崩潰（mode collapse）。
適用於少量數據：對比學習能有效提高少量標註字型數據的泛化能力。

這樣應該能讓你的字型風格轉換效果更好！🚀

Gemini

好的，為了在您的 Zi2ZiModel 中加入 Contrastive Loss，我們需要進行以下修改：

添加特徵提取器： 在生成器中添加一個特徵提取器，用於提取風格特徵。
實現 Contrastive Loss： 實現 Contrastive Loss 函數。
修改 forward 函數： 在 forward 函數中，使用特徵提取器提取風格特徵，並計算 Contrastive Loss。
修改 backward_G 函數： 在 backward_G 函數中，添加 Contrastive Loss 的計算。
修改 optimize_parameters 函數： 在 optimize_parameters 函數中，添加 Contrastive Loss 的優化。

以下是修改後的程式碼：

Python

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.nn.utils import spectral_norm
from torch.utils.data import DataLoader
import os
import time
import math
import functools
from torch.nn import init

# 假設 UnetSkipConnectionBlock, DatasetFromObj 已定義

class FeatureExtractor(nn.Module):
    def __init__(self, input_nc=1, output_nc=128, ngf=64):
        super(FeatureExtractor, self).__init__()
        self.model = nn.Sequential(
            spectral_norm(nn.Conv2d(input_nc, ngf, kernel_size=4, stride=2, padding=1)),
            nn.LeakyReLU(0.2, True),
            spectral_norm(nn.Conv2d(ngf, ngf * 2, kernel_size=4, stride=2, padding=1)),
            nn.LeakyReLU(0.2, True),
            spectral_norm(nn.Conv2d(ngf * 2, ngf * 4, kernel_size=4, stride=2, padding=1)),
            nn.LeakyReLU(0.2, True),
            spectral_norm(nn.Conv2d(ngf * 4, output_nc, kernel_size=4, stride=1, padding=0)),
            nn.Flatten()
        )

    def forward(self, x):
        return self.model(x)

class ContrastiveLoss(nn.Module):
    def __init__(self, margin=1.0):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin

    def forward(self, output1, output2, label):
        euclidean_distance = F.pairwise_distance(output1, output2)
        loss_contrastive = torch.mean((1 - label) * torch.pow(euclidean_distance, 2) +
                                      (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
        return loss_contrastive

class UNetGenerator(nn.Module):
    # ... (您的 UNetGenerator 程式碼)
    pass

class Discriminator(nn.Module):
    # ... (您的 Discriminator 程式碼)
    pass

class Zi2ZiModel:
    def __init__(self, input_nc=1, output_nc=1, ngf=64, ndf=64, lr=0.001, gpu_ids=None, save_dir='.', is_training=True, image_size=256, epochs=100, margin=1.0):
        # ... (您的 __init__ 程式碼)
        self.margin = margin
        self.feature_extractor = FeatureExtractor(input_nc=input_nc, output_nc=128, ngf=ngf)
        self.contrastive_loss = ContrastiveLoss(margin=margin)

    # ... (您的 init_weights, init_net, setup 程式碼)

    def setup(self):
        super().setup() # call super class setup
        self.feature_extractor = FeatureExtractor(input_nc=self.input_nc, output_nc=128, ngf=self.ngf)

        self.init_net(self.feature_extractor, gpu_ids=self.gpu_ids)

        self.optimizer_F = optim.Adam(self.feature_extractor.parameters(), lr=self.lr, betas=(0.5, 0.999))
        self.scheduler_F = CosineAnnealingLR(self.optimizer_F, T_max=self.epochs, eta_min=0)

    def set_input(self, real_A, real_B):
        # ... (您的 set_input 程式碼)

    def forward(self):
        self.fake_B = self.netG(self.real_A)
        self.real_features = self.feature_extractor(self.real_B)
        self.fake_features = self.feature_extractor(self.fake_B)

    def backward_D(self):
        # ... (您的 backward_D 程式碼)

    def backward_G(self):
        fake_AB = torch.cat([self.real_A, self.fake_B], 1)
        loss_G_GAN = self.criterion(self.netD(fake_AB), torch.ones_like(self.netD(fake_AB)))
        contrastive_loss = self.contrastive_loss(self.real_features, self.fake_features, torch.zeros(self.real_features.size(0)).to(self.real_features.device)) # label 0 for similar pairs

        loss_G = loss_G_GAN + contrastive_loss
        loss_G.backward()
        return loss_G_GAN, contrastive_loss

    def optimize_parameters(self):
        self.forward()
        self.optimizer_D.zero_grad()
        loss_D = self.backward_D()
        self.optimizer_D.step()

        self.optimizer_G.zero_grad()
        self.optimizer_F.zero_grad()
        loss_G_GAN, contrastive_loss = self.backward_G()
        self.optimizer_G.step()
        self.optimizer_F.step()

        return loss_D, loss_G_GAN, contrastive_loss

    def update_lr(self):
        # ... (您的 update_lr 程式碼)
        self.scheduler_F.step()
        print(f"Learning rate updated: F={self.optimizer_F.param_groups[0]['lr']:.6f}")

# ... (您的訓練迴圈程式碼)

主要修改：

FeatureExtractor 類別： 添加了一個特徵提取器，用於提取風格特徵。
ContrastiveLoss 類別： 實現了 Contrastive Loss 函數。
forward 函數： 使用特徵提取器提取風格特徵，並將其存儲在 self.real_features 和 self.fake_features 中。
backward_G 函數： 添加了 Contrastive Loss 的計算，並將其添加到總的生成器損失中。
optimize_parameters 函數： 添加了 self.optimizer_F 的優化步驟。
setup 函數： 對 feature_extractor 進行初始化，並添加 optimizer_F
update_lr 函數： 添加了 scheduler_F 的更新步驟。

請注意，您可能需要根據您的具體任務和數據集調整 Contrastive Loss 的參數（例如，margin）。

Max的程式語言筆記