现在来利用 pytorch 提供的高级功能来自己实现一个 MLP 用来处理图像分类任务。
加载数据集。
from d2l import torch as d2l batch_size = 256 train_iter,test_iter = d2l.load_data_fashion_mnist(batch_size)
|
定义模型。
net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
|
初始化权重。这一步将会遍历模型的每一层,并为每个线性层(type(m) == nn.Linear)随机初始化权重(默认均值为 0 ,标准差为 0.01)。
def init_weights(m): if type(m) == nn.Linear: nn.init.normal_(m.weight, std=0.01) net.apply(init_weights)
|
定义一个交叉熵 loss 函数。
loss = nn.CrossEntropyLoss(reduction='none')
|
定义一个优化器 trainer( updater ),学习率为 0.1 ,使用 小批量随机梯度下降:
trainer = torch.optim.SGD(net.parameters(), lr=0.1)
|
训练
def train_epoch(net,train_iter,loss,updater): net.train() metric = Accumulator(3) for X,y in train_iter: y_hat = net(X) l = loss(y_hat,y) updater.zero_grad() l.mean().backward() updater.step() metric.add(float(l.sum()), accuracy(y_hat, y), y.numel()) return (metric[0]/metric[2]),(metric[1]/metric[2]) def train(net,train_iter,test_iter,loss,num_epochs,updater): for epoch in range(num_epochs): train_metrics = train_epoch(net,train_iter,loss,updater) test_acc = evaluate_accuracy(net, test_iter) print(f"Epoch {epoch} , Acc = {test_acc}") train_loss, train_acc = train_metrics
|
训练 30 轮,在这个任务上足够了:
num_epochs = 30 train(net,train_iter,test_iter,loss,num_epochs,trainer)
|
执行预测:
def predict(net,test_iter,n=6): for X, y in test_iter: break trues = d2l.get_fashion_mnist_labels(y) preds = d2l.get_fashion_mnist_labels(net(X).argmax(axis=1)) for tr,pred in zip(trues,preds): print(f"True: {tr} , Predict: {pred}") predict(net, test_iter)
|