softmax 回归的简洁实现

导入所需包和模块

import torch
from torch import nn
from torch.nn import init # 初始化模块
import numpy as np
import sys
sys.path.append("..")
import d2lzh_pytorch as d2l

定义模型

# 定义超参数
num_inputs, num_outputs, num_hiddens = 784, 10, 256
net = nn.Sequential(d2l.FlattenLayer(),  # 改变X的形状
                   nn.Linear(num_inputs, num_hiddens),
                   nn.ReLU(), # 使用relu激活函数
                   nn.Linear(num_hiddens, num_outputs),
                   )

for params in net.parameters():
    init.normal_(params, mean=0, std=0.01)  # 初始化参数

读取数据，训练模型

batch_size = 256  # 设定批量大小为256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
loss = torch.nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(net.parameters(), lr = 0.5)

num_epochs = 5  # 迭代周期为5
d2l.train_ch3(net, train_iter, test_iter, loss,
             num_epochs, batch_size,None, None, optimizer)