典型的动漫图像风格可以用一个6维向量来描述

社区文章发布于 2025 年 8 月 2 日

Yuyao Huang

Fgdfgfthgr

在对来自 Danbooru 数据集的动漫图像进行图像嵌入模型训练后，我注意到一些有趣的事情：输出的内在维度始终在 6 左右！

这意味着什么？这意味着你只需用 6 个数字就可以完全描述典型动漫图像的风格！

在本文中，我将介绍

它为什么有用
训练配置
风格如何分布
这 6 个数字在图像中各自代表什么

这有什么用？我为什么要这样做？

最初，我的目标是从头开始训练一个小型扩散模型，使用 Danbooru 数据集，以熟悉各种扩散模型概念。

然而，许多扩散模型选择使用艺术家标签来控制输出图像的风格。这在基于动漫图像的模型中尤其常见，比如许多 SDXL 微调模型。

我真的不喜欢那样，原因有三：

许多艺术家的风格非常相似，导致许多艺术家标签冗余。
有些艺术家的作品中不止有一种独特的艺术风格。例如，草图与完成的图像。
容易出现内容渗漏。如果你选择的艺术家标签包含大量重复内容，这些内容很可能会渗漏到你的输出中，即使你没有提示它们。

当然，对于第三个原因，你可以通过使用负面提示来避免。但有时即使那样也不够有效。如果风格能够更直接地控制，那就更容易了。

因此，主要受 PonyV7 模型的启发，我产生了使用风格嵌入模型的想法。

这意味着什么？这意味着一个模型可以接收任意大小的图像，并为每张图像输出一个风格向量。风格向量存在于 N 维空间中，本质上只是一个长度为 N 的数字列表。列表中的每个数字对应于输入图像所具有的特定风格元素。

如果我们能以某种方式将这个风格向量与训练图像一起注入到扩散模型中，理论上我们可以获得一个使用这个风格向量来确定输出风格的扩散模型，而不是依赖艺术家标签。

有哪些风格元素，有多少？我还不清楚，这要留待后续训练来解决。但可能像粗细线条、写实与卡通、凌乱与整洁等等。

所以我深入研究了 HuggingFace，试图找出是否有人有过类似的想法。令人惊讶的是，我找不到任何针对动漫图像执行此操作的模型。

这意味着我必须构建自己的模型。

大约在同一时间，事实证明，训练一个质量可接受的扩散模型对我拥有的显卡来说太过分了。所以我放弃了这个想法，专注于风格嵌入模型。希望它对其他人有用，比如现在正在阅读这篇文章的你！

训练配置

训练这个风格嵌入模型很难！它比分类任务难得多，而且非常容易过拟合。难怪我找不到其他人在做这个。

经过无数次尝试，我终于找到了一个配置，为这项任务生成了一个拟合良好、体积小的模型。

数据集

准备数据集可能是许多深度学习任务中最耗时的一步。最初我尝试使用简单的对比学习对，然后是三元组（https://lilianweng.github.io/posts/2021-05-31-contrastive/）。它们的大小为 6-8k，即使我使用了所有的数据增强，它们要么不起作用，要么出现严重的过拟合。

三元组：[锚点、正例、负例]。正例是与锚点图像风格相似的图像，而负例是与锚点图像风格不相似的图像

一些研究表明，我可能需要 10k+ 三元组才能进行有效学习。我绝对不想手动标注那么多。

然而，我们使用的是 Danbooru 数据集。它有两个我们可以利用的优点：

同一艺术家的图像可能共享相似的风格。
不同艺术家的图像风格更可能不相似。

利用这些特性，我提出了一种解决方案，可以根据作者级别的注释动态构建新的三元组。这大大减少了手动注释的工作量。

训练真实数据是通过这种两部分方法获得的。第一部分：

查找至少有 5 张图像但最多有 60 张图像的艺术家。
按字母顺序排列。
逐一检查他们的作品集。如果艺术家具备独特且稳定的风格，则手动删除与该艺术家其他图像风格过于不相似的图像。然后将该艺术家及其所有剩余作品记录到一个文本文件中。

通过这样做，我得到了一个如下所示的文本文件：

artist1: 1.webp 2.webp 3.webp 4.webp 5.webp
artist2: 6.webp 7.webp 8.webp 9.webp 10.webp
...

每行代表一组我知道彼此风格相似的图像。

我跳过了没有稳定和独特风格的艺术家。我还跳过了一旦我已经记录了十几位艺术家后，那些风格过于常见的艺术家。

517位艺术家被记录为我的训练集。（不包括跳过的艺术家）

第 2 部分（使用第 1 部分的数据）

选择一个不满足条件 A 的随机艺术家 (a)。然后选择另一个随机艺术家 (b)。
比较他们的作品集，以确定他们的风格是“相似”、“不相似”还是“我无法判断”。
(a) 保持不变，用新的随机艺术家替换 (b)。
重复步骤 2 和 3，直到 (a) 满足条件 A。
重复步骤 1 到 4，直到所有艺术家都满足条件 A。

条件 A：该艺术家至少有 4 位其他艺术家被标记为与其“不相似”。

通过这样做，我们获得了一个如下所示的文本文件：

artist1 vs artist2: 1
artist3 vs artist4: 0
artist5 vs artist6: NA
...

1 表示相似，0 表示不相似，NA 表示我无法判断。

共收集了 2406 对相似/不相似对。（其中大部分是不相似的）

在标注过程中，我尽量保持客观，这意味着我只在我非常确定它们具有相似风格时才将它们标注为“相似”，或者在我非常确定它们不同时才标注为“不相似”。

然而，这不会是“完美”的。因为风格是一个复杂的事物，每个人对其定义都不同。

我创建的这两个文本文件可以在这里获取。

总之，有了上面创建的文本文件，我们就可以使用以下过程动态生成三元组：

选择一个随机图像作为锚点图像
找出该图像所属的艺术家
从风格相似的艺术家中选择随机图像，形成正例图像
从风格不相似的艺术家中选择随机图像，形成负例图像

请注意，一个艺术家的风格总是与他自己相似。

验证集呢？我在早期的方法中手动标注了一些三元组。现在它们被用作验证集。

预处理

由于网络能够接受任意大小的输入，因此无需填充。我只使用了相当简单的图像增强：

import torch.nn as nn
from torchvision.transforms import v2

def closest_interval(img, interval=8):
    c, h, w = img.shape
    new_h = h - (h % interval) if h % interval != 0 else h
    new_w = w - (w % interval) if w % interval != 0 else w
    h_start = (h - new_h) // 2
    w_start = (w - new_w) // 2
    new_h, new_w = max(new_h, interval), max(new_w, interval)
    return img[:, h_start:h_start + new_h, w_start:w_start + new_w]


class RandomSizeTransform(nn.Module):
    def __init__(self, smallest_ratio, size_range):
        super(RandomSizeTransform, self).__init__()
        self.smallest_ratio = smallest_ratio
        self.size_range = size_range

    def forward(self, img):
        c, h, w = img.shape
        ratio = random.uniform(self.smallest_ratio, 1)
        target_h, target_w = int(h * ratio), int(w * ratio)
        h_start, w_start = random.randint(0, h-target_h), random.randint(0, w-target_w)
        img = img[:, h_start:h_start+target_h, w_start:w_start+target_w]
        target_size = random.randint(*self.size_range)
        img = adj_size(img, size=target_size)
        return closest_interval(img)
        
transforms = v2.Compose([
    RandomSizeTransform(0.8, (1024, 1400)),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomVerticalFlip(p=0.5),
    v2.ColorJitter(0.3, 0.3, 0.2, 0.2),
    v2.RandomGrayscale(p=0.2),
])

转换后，图像的像素值被缩放到 -1 到 +1 之间：`img = 2*(img/255)-1`

请注意，我在这里使用了随机灰度化。我个人认为图像的灰度化版本与原始图像具有相同的风格。如果您正在训练自己的图像嵌入模型，您可能希望禁用此功能。

损失函数

我的目标是让风格相似的嵌入彼此靠近，而风格不相似的嵌入彼此远离。

显而易见的选择是三元组损失，它试图使 `neg_dis` 至少比 `pos_dis` 大 K：`loss = max(K - (neg_dis - pos_dis), 0)`。`pos_dis` 是锚点样本与正样本之间的距离，而 `neg_dis` 是锚点样本与负样本之间的距离。

K 是边距距离，可以是任何数字，我在这里选择了 3。

然而，普通的三元组损失在这里效率很低。因为我们不做任何填充，并且我们的模型具有可变输入大小，所以我们只能使用大小为 1 的批次。因此，每个训练步骤网络只能看到 3 张图像和 1 个三元组。提高训练效率的一种明显方法是在每个训练步骤中利用更多样本。

所以我设计了

class MultiSampleTripletLossAllPairs(nn.Module):
    def __init__(self, margin=1.0, p=2):
        """
        Considers all positive-negative pairs

        Args:
            margin: Margin for triplet loss
            p: The norm degree for pairwise distance
            reduction: 'mean', 'sum' or 'none'
        """
        super(MultiSampleTripletLossAllPairs, self).__init__()
        self.margin = margin
        self.p = p
       
    
    def calculate_average_norms(self, positive, negative):
        pos_norms = torch.linalg.vector_norm(positive, ord=self.p, dim=1)  # (1, N)
        neg_norms = torch.linalg.vector_norm(negative, ord=self.p, dim=1)  # (1, M)

        return pos_norms.mean(), neg_norms.mean()

    def forward(self, anchor, positive, negative):
        """
        Args:
            anchor: (1, D) - single anchor embedding
            positive: (N, D) - N positive samples
            negative: (M, D) - M negative samples
        Returns:
            loss value
        """
        avg_pos_norm, avg_neg_norm = self.calculate_average_norms(positive, negative)
        anchor, positive, negative = anchor.unsqueeze(0), positive.unsqueeze(0), negative.unsqueeze(0)
        # Compute positive distances (1, 1, N)
        pos_dist = torch.cdist(anchor, positive, p=self.p)
        # Compute negative distances (1, 1, M)
        neg_dist = torch.cdist(anchor, negative, p=self.p)

        # Compute all possible triplet combinations (N x M)
        pos_dist = pos_dist.permute(2, 1, 0).squeeze(1)  # (N, 1)
        neg_dist = neg_dist.squeeze(0)  # (1, M)

        # Compute loss for all combinations
        losses = F.relu(pos_dist - neg_dist + self.margin)

        # Apply reduction
        return losses.mean(), pos_dist.mean(), neg_dist.mean(), avg_pos_norm, avg_neg_norm

对于 N 个正嵌入和 M 个负嵌入，它相当于 N*M 个三元组。这可以平滑损失曲线并加速训练过程。

为了防止嵌入彼此过于遥远，我还将范数作为惩罚项添加。`loss += torch.linalg.vector_norm(anchor_embd, ord=2, dim=1).mean() * 0.001`

网络架构

最初，我尝试了一种类似于 VGG 的架构。图像经过一系列卷积和池化层，然后通过 AdaptiveAvgPool2d 缩减到特定大小，再展平并通过 MLP 获得最终嵌入。

效果不太好。

后来我找到了Gram 矩阵。它使用矩阵计算来获取每个通道之间的相似性，然后将这些相似性结果输入到 MLP 中。我用 Gram 矩阵替换了自适应池化。

它有一些不错的特性：

空间不变性
输出具有固定形状
理论上，它非常擅长捕捉图像的风格

class CompactGramMatrix(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels
        # Precompute indices for lower triangle (including diagonal)
        self.register_buffer('tril_indices',
                             torch.tril_indices(in_channels, in_channels, offset=0, dtype=torch.int32))

    def forward(self, x):
        """
        Input: (B, C, H, W)
        Output: (B, C*(C+1)//2) compact Gram features
        """
        b, c, h, w = x.size()
        x = x.view(b, c, -1) / ((h * w) ** 0.5)  # Flatten spatial dimensions -> (B, C, H*W), then normalise

        # Compute full Gram matrix (still needed temporarily)
        gram = torch.bmm(x, x.transpose(1, 2))  # (B, C, C)

        # Extract lower triangle including diagonal
        compact_gram = gram[:, self.tril_indices[0], self.tril_indices[1]]  # (B, n_unique)
        return compact_gram

注意我只使用了下三角，因为 Gram 矩阵的上三角包含与下三角完全相同的信息。

我的模型看起来像这样（查看Hugging Face 模型页面了解更多详情！）

class EmbeddingNetwork(nn.Module):
    def __init__(self):
        super(EmbeddingNetwork, self).__init__()
        self.input_conv = nn.Conv2d(3, 16, 5, padding='same', padding_mode='reflect', bias=False)
        self.conv1 = ResBlock(16, 3, 2)
        self.pool1 = ConvPool(16, 32) # 2
        self.conv2 = ResBlock(32, 3, 2)
        self.pool2 = ConvPool(32, 64) # 4
        ...
        self.gram = CompactGramMatrix(256)
        self.compact = nn.Linear(256*(256+1)//2, 1024, bias=False)
        self.conpactnorm = nn.LayerNorm(1024, elementwise_affine=False)
        self.fc1 = nn.Linear(1024, 1024, bias=False)
        self.fc1norm = nn.LayerNorm(1024, elementwise_affine=False)
        self.act = nn.LeakyReLU(inplace=True)
        ...

    def forward(self, x):
        x = self.input_conv(x)
        x = self.pool1(self.conv1(x))
        x = self.pool2(self.conv2(x))
        ...
        x = self.gram(x)
        x = self.compact(x)
        x = self.conpactnorm(x)
        x = self.act(self.fc1norm(self.fc1(x)))
        ...

训练超参数

训练是使用 PyTorch Lightning 完成的。

学习率 = 0.0001

权重衰减 = 0.0001

AdEMAMix 优化器

ExponentialLR 调度器，伽马值为 0.99，每个 epoch 应用一次。

批处理大小为 1。accumulate_grad_batches 为 16。

每个锚点图像使用 16 个正样本和 16 个负样本。

在 2 个 A100 GPU 上训练了 15 个 epoch。总共进行了 3434 次优化器更新。

发现：风格分布

我用完全相同的配置训练了 5 个模型，但输出维度不同：128、32、16、8 和 6。（128 是首先训练的）

128 维对于描述图像风格来说应该太多了。然而，通过使用 skdim 包进行内在维度估计，我们可以大致了解最小输出维度会是怎样的。

import skdim

estimators = [skdim.id.TwoNN(), skdim.id.CorrInt(), skdim.id.DANCo()]
results = {}

for est in estimators:
    est.fit(predictions)
    results[type(est).__name__] = est.dimension_

print("Intrinsic Dimension Estimates:")
for name, dim in results.items():
    print(f"{name}: {dim:.2f}")

每次评估期间使用了 5000 张来自整个数据集的随机图像。我在这里的维度数字在不同运行之间是稳定的，在 +-0.02 之间。

结果