Pytorch+MINST 全连接神经网路训练学习

声明

学习地址：https://www.youtube.com/watch?v=wnK3uWv_WkU&list=PLhhyoLH6IjfxeoooqP9rhU3HJIAVAJ3Vz&index=4

代码

导入模块

import torch
import torch.nn as nn  # torch的网路库，本例程用到nn.Linear
import torch.optim as optim  #
import torch.nn.functional as F  # torch 函数库
from torch.utils.data import DataLoader  #DataLoder
import torchvision.datasets as datasets  # 数据集
import torchvision.transforms as transforms  # 图像变换库

Creat Fully Connected Network

class NN(nn.Module):
    def __init__(self, input_size, num_classes):
        super(NN, self).__init__()
        self.fc1 = nn.Linear(input_size, 50)
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

super方法的使用

参考文章

Set device

我GPU不行，注释掉不用

# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')