使用transfomrer的编码器实现minist数字识别

原创

已于 2024-08-08 10:25:22 修改 · 626 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python #图像处理 #目标跟踪

于 2024-08-07 18:41:04 首次发布

使用transfomrer的编码器实现minist数字识别

1.数据预处理
2. 定义transformer模型
3. 训练
4.测试
5.测试单张图片的可视化

1.数据预处理

import torch
import torch.nn as nn
import torch.optim as optim

from torchvision import datasets,transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

2. 定义transformer模型

class TransformerModel(nn.Module):
    def __init__(self, input_dim, num_classes,n_heads=4,num_encoder_layers=3):
        super(TransformerModel, self).__init__()
        self.positional_encoding=nn.Parameter(torch.zeros(1, 28, 28))
        encoder_layers=self.transformer=nn.TransformerEncoderLayer(d_model=28,nhead=n_heads)
        self.transformer_encoder = nn.T