PyTorch深度学习实战——多分类问题

最新推荐文章于 2024-07-24 14:59:26 发布

CheneyTAT

最新推荐文章于 2024-07-24 14:59:26 发布

阅读量537

点赞数

分类专栏： PyTorch深度学习实践文章标签： pytorch 深度学习

本文链接：https://blog.csdn.net/weixin_66648679/article/details/124073688

版权

PyTorch深度学习实践专栏收录该内容

7 篇文章 1 订阅

订阅专栏

import torch
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim

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

train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, batch_size=batch_size)
train_loader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, transform=transform)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        return self.l5(x)  # 最后一层不需要用relu函数处理，因为最后一层输出的处理被包含进了

代码存档，还未学完，暂时用不上