pytorch中如何在vit_b_16后加入全连接层

在 PyTorch 中，可以使用 nn.Sequential() 函数来将 ViT-B/16 模型和全连接层组合在一起。以下是一个示例代码：

import torch.nn as nn
from transformers import ViTModel

# 加载 ViT-B/16 模型
vit_model = ViTModel.from_pretrained('google/vit-base-patch16-224')

# 定义全连接层
fc_layer = nn.Linear(in_features=768, out_features=10)

# 将 ViT-B/16 模型和全连接层组合在一起
model = nn.Sequential(vit_model, fc_layer)

在这个示例中，我们首先使用 `ViTModel.from_pretrained()` 函数加载了 ViT-B/16 模型。然后，我们定义了一个全连接层，其中输入特征数为 768（ViT-B/16 模型的输出特征数），输出特征数为 10。最后，我们使用 `nn.Sequential()` 函数将 ViT-B/16 模型和全连接层组合在一起，形成一个新的模型 `model`。

需要注意的是，这个示例中的全连接层只是一个简单的示例，实际应用中需要根据具体任务来设计合适的全连接层。