Class10简洁实现

Class10简洁实现

import torch
from torch import nn
from d2l import torch as d2l

# 输入为28*28，输出为10类，第1、2隐藏层256神经元
num_inputs, num_outputs, num_hiddens1, num_hiddens2 = 784, 10, 256, 256
# 第1个隐藏层丢弃率为0.2，第2个为0.5
dropout1, dropout2 = 0.2, 0.5

# nn.Flatten():把28*28展平为784
net = nn.Sequential(nn.Flatten(),# 输入层->第1隐藏层nn.Linear(784, 256),# ReLU激活nn.ReLU(),# 在第一个全连接层之后添加一个dropout层nn.Dropout(dropout1),# 第1隐藏层->第2隐藏层nn.Linear(256, 256),# ReLU激活nn.ReLU(),# 在第二个全连接层之后添加一个dropout层nn.Dropout(dropout2),# 第2隐藏层->输出10类nn.Linear(256, 10))# 初始化权重函数
def init_weights(m):# 判断如果为线性if type(m) == nn.Linear:# 正态分布初始化，均值为0，标准差为0.01nn.init.normal_(m.weight, std=0.01)
# 若为nn.Linear，则调用init_weight函数进行初始化
net.apply(init_weights);

# 设置训练轮数，学习率，批次大小
num_epochs,lr,batch_size = 10,0.5,256
# 定义损失函数，并保留每个样本损失
loss = nn.CrossEntropyLoss(reduction='none')
# 加载训练集和测试集
train_iter,test_iter = d2l.load_data_fashion_mnist(batch_size)
# 设置SGD随机梯度下降优化器
trainer = torch.optim.SGD(net.parameters(),lr=lr)
# 调用训练主函数
d2l.train_ch3(net,train_iter,test_iter,loss,num_epochs,trainer)