多模态基础--CNN和ViT

# 位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length):
        super().__init__()
        # 类别token
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))
        # 初始化编码矩阵 (max_len, d_model)
        pe = torch.zeros(size=(max_seq_length, d_model))
        # 当前词在序列中的位置 (max_len, 1)
        pos = torch.arange(0, max_seq_length).unsqueeze(1)
        # 表示公式中2i (d_model/2, )
        _2i = torch.arange(0, d_model, 2)
        # 计算10000**(2i/d_model) (d_model/2, )
        div_term = torch.pow(10000, (_2i / d_model))
        # 按奇偶数维度计算位置编码值 (max_len, d_model)
        pe[:, 0::2] = torch.sin(pos / div_term)
        if d_model % 2 == 1:
            _2i1 = torch.arange(0, d_model-1, 2)
            div_term = torch.pow(10000, (_2i1 / d_model))
        pe[:, 1::2] = torch.cos(pos / div_term)
        self.register_buffer("pe", pe)

    def forward(self, x):
        # 为批次中的每张图片分配一个类别token
        tokens_batch = self.cls_token.expand(x.size()[0], -1, -1)
        # 将类别token添加到每个图像的补丁嵌入数组的开头
        x = torch.cat((tokens_batch, x), dim=1)
        # 将位置编码添加到嵌入中
        x = x + self.pe
        return x

class VisionTransformer(nn.Module):
    def __init__(self,d_model,n_classes,img_size,patch_size,n_channels,n_heads,n_layers):
        super().__init__()

        assert img_size[0] % patch_size[0] == 0 and img_size[1] % patch_size[1] == 0, "img_size 必须能被 patch_size 整除"
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"

        self.d_model = d_model  # 模型维度，嵌入的维度（宽度）
        self.n_classes = n_classes  # 类别的数量
        self.img_size = img_size  # 图片大小
        self.patch_size = patch_size  # 补丁大小
        self.n_channels = n_channels  # 通道数
        self.n_heads = n_heads  # 注意力头的数量
        # 补丁的数量 = (32x32) // (4x4)
        self.n_patches = (self.img_size[0] * self.img_size[1]) // (self.patch_size[0] * self.patch_size[1])
        # 序列的长度 = 1（分类token） + 补丁的数量
        self.max_seq_length = self.n_patches + 1
        # 补丁嵌入
        self.patch_embedding = PatchEmbedding(
            self.d_model,
            self.patch_size,
            self.n_channels
        )
        # 位置编码
        self.positional_encoding = PositionalEncoding(
            self.d_model,
            self.max_seq_length
        )
        self.transformer_encoder = nn.Sequential(*[
            TransformerEncoder(self.d_model, self.n_heads)
            for _ in range(n_layers)
        ])

        # 用于分类的MLP
        self.classifier = nn.Sequential(
            nn.Linear(self.d_model, self.n_classes),
            nn.Softmax(dim=-1)
        )

    def forward(self, images):
        # 将图片转换成补丁的嵌入（embedding）
        x = self.patch_embedding(images)
        # 添加位置编码
        x = self.positional_encoding(x)
        # 编码
        x = self.transformer_encoder(x)
        # 分类的线性层
        x = self.classifier(x[:, 0])
        return x


# 基础配置
ROOT_DIR = Path(__file__).parent.parent
device = 'cuda' if torch.cuda.is_available() else 'cpu'
log_dir = ROOT_DIR / 'logs'

# 超参数配置
d_model = 12  # 嵌入的维度12
n_classes = 10  # 类别数量为10
img_size = (32, 32)  # 图片大小为32x32
patch_size = (16, 16)  # 补丁的大小是16x16
n_channels = 1  # 灰度图片通道数量为1
n_heads = 4  # 4个注意力头
n_layers = 3  # 3层编码器
batch_size = 128  # 每个批次128张图片
epochs = 10  # 训练10个epoch
alpha = 0.005  # 学习率5e-3

# 准备数据
transform = T.Compose([T.Resize(img_size),T.ToTensor()])
train_set = MNIST(
    root="./datasets", train=True, download=True, transform=transform
)
test_set = MNIST(
    root="./datasets", train=False, download=True, transform=transform
)

train_loader = DataLoader(train_set, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_set, shuffle=False, batch_size=batch_size)

# 初始化模型
ViT = VisionTransformer(
    d_model,
    n_classes,
    img_size,
    patch_size,
    n_channels,
    n_heads,
    n_layers
).to(device)

# 初始化优化器
optimizer = Adam(ViT.parameters(), lr=alpha)
# 初始化损失函数
loss = nn.CrossEntropyLoss()

# 开始训练
with SummaryWriter(log_dir=str(log_dir / time.strftime('%Y-%m-%d_%H-%M-%S'))) as writer:
    for epoch in range(epochs):
        training_loss = 0.0
        for i, data in enumerate(train_loader, 0):
            # 取出图像和对应的标签
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = ViT(inputs)
            # 交叉熵损失
            loss_v = loss(outputs, labels)
            loss_v.backward()
            optimizer.step()
            training_loss += loss_v.item()
        writer.add_scalar('loss', training_loss, epoch + 1)

        print(
            f'Epoch {epoch + 1}/{epochs} loss: {training_loss / len(train_loader):.3f}')

correct = 0
total = 0
# 开始测试
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        outputs = ViT(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print(f'\n预测准确率: {100 * correct // total} %')