#!/usr/bin/env python
# coding: utf-8

# # softmax回归的简洁实现
# 
# 我们在[“线性回归的简洁实现”](linear-regression-gluon.ipynb)一节中已经了解了使用PyTorch实现模型的便利。下面，让我们再次使用PyTorch来实现一个softmax回归模型。首先导入所需的包或模块。

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import d2ltorch as d2lt
import torch
from torch import nn, optim


# ## 获取和读取数据
# 
# 我们仍然使用Fashion-MNIST数据集和上一节中设置的批量大小。

# In[2]:


root = '~/dataset/'
batch_size = 256
train_iter, test_iter = d2lt.load_data_fashion_mnist(root, batch_size=batch_size)


# ## 定义和初始化模型
# 
# 在[“softmax回归”](softmax-regression.ipynb)一节中提到，softmax回归的输出层是一个全连接层。因此，我们添加一个输出个数为10的全连接层。我们使用均值为0、标准差为0.01的正态分布随机初始化模型的权重参数。
# 
# *这里我们提前使用 [“4.2. 模型参数的访问、初始化和共享”](../chapter_deep-learning-computation/parameters.ipynb) 一节中定义的初始化方法 `params_init(model, init, **kwargs)`*

# In[3]:


class MySoftMax(nn.Module):
    def __init__(self, **kwargs):
        super(MySoftMax, self).__init__(**kwargs)
        self.linear = nn.Linear(28 * 28, 10)
    
    def forward(self, x):
        x = x.reshape(-1, 28 * 28)
        return self.linear(x)

net = MySoftMax()
d2lt.params_init(model=net, init=nn.init.normal_, std=0.01)


# ## softmax和交叉熵损失函数
# 
# 如果做了上一节的练习，那么你可能意识到了分开定义softmax运算和交叉熵损失函数可能会造成数值不稳定。因此，PyTorch提供了一个包括softmax运算和交叉熵损失计算的函数。它的数值稳定性更好。

# In[4]:


loss = nn.CrossEntropyLoss()


# ## 定义优化算法
# 
# 我们使用学习率为0.1的小批量随机梯度下降作为优化算法。

# In[5]:


optimizer = optim.SGD(net.parameters(), lr=0.1)


# ## 训练模型
# 
# 接下来，我们使用上一节中定义的训练函数来训练模型。

# In[6]:


num_epochs = 5
d2lt.train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, None, None, optimizer)


# ## 小结
# 
# * PyTorch提供的函数往往具有更好的数值稳定性。
# * 可以使用PyTorch更简洁地实现softmax回归。
# 
# ## 练习
# 
# * 尝试调一调超参数，如批量大小、迭代周期和学习率，看看结果会怎样。
# 
# 
# 
# ## 扫码直达[讨论区](https://discuss.gluon.ai/t/topic/740)
# 
# ![](../img/qr_softmax-regression-gluon.svg)