pytorch的MINST数据集示例

最新推荐文章于 2024-07-19 16:12:05 发布

置顶阆遤

最新推荐文章于 2024-07-19 16:12:05 发布

阅读量151

点赞数 2

文章标签： pytorch 深度学习人工智能 python

本文链接：https://blog.csdn.net/liangma/article/details/140509335

版权

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms

# 定义一个简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 初始化模型、优化器和损失函数
model = SimpleCNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 训练模型
epochs = 5
for epoch in range(epochs):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print(f'Epoch [{epoch+1}/{epochs}] Batch [{batch_idx}/{len(train_loader)}] Loss: {loss.item()}')

print("训练完成")

上述代码我的目的是测试cuda+cudnn的安装环境是否正确。在下载MNIST时，出现无法下载部分文件的问题。经过寻找，将MNIST数据打包放在我的资源中（https://download.csdn.net/download/liangma/89550954），并将由fittencode辅助生成的示例代码发出来，以利有需要的人参考使用。

注意事项：安装好pytorch和cuda+cudnn，将上述代码保存为py文件，在文件当前目录下建立data文件夹，MNIST数据集下载下来，把4个gz压缩文件放到.\data\MNIST\raw目录下,运行上述代码，如果提示需要下载某个gz文件，进入.\data\MNINST\raw目录，将gz文件解压所在其中就可。