用于预训练视觉变压器的自定义补丁嵌入层

Question

我正在寻找一种为 vanilla ViT 使用自定义补丁嵌入层的方法。我想要预训练模型中的其余 ViT。有没有办法使用 Pytorch 来做到这一点？

我可以加载预先训练的模型，也可以从头开始编写 ViT 代码。但我想要一些可以在补丁嵌入后使用模型权重的东西。

谢谢，

Answer 1

这在很大程度上取决于两个模型的架构，但你可以这样做：

class MyVIT(nn.Module):
    def __init__(self, embedding, vit_model):
        super().__init__()
        self.embedding = embedding
        self.vit_model = vit_model

    def forward(self, x):
        x = self.embedding(x)
        x = self.vit_model(x)
        return x

在此示例中，

embedding

将是您的自定义嵌入，

vit_model

将是除嵌入之外的 vit 模型的所有训练层。根据 vit 模型的结构方式，您可能需要侵入它以提取非嵌入层，以便您可以简单地将输入传递给它们。