PyTorch中有一个torch.nn提供了所有你在组建神经网络过程中需要的模块。每个神经网络模型都是PyTorch的nn.Module的子类，神经网络是由层或者模型组成的，也就是说一个神经网络是其他层或者网络组合起来的。因为神经网络的这种层级式结构，也让我们很容易就能建立好复杂的网络模型。

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

首先就是导包，这里没什么可解释的了。

使用GPU加速训练

如果我们有GPU，我们可以使用GPU加速模型训练。可以使用torch.cuda检查硬件是否有GPU，如果有则使用GPU，没有则继续使用CPU训练。

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

如果有GPU，cuda可用，会显示：

Using cuda device

如果你没有GPU那就用CPU也没什么关系。

创建网络

为nn.Module创建子类，定义我们的神经网络。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

__init__：在这里创建网络结构。
- self.flatten我们接受的输入是一个图片的张量，这里要将其展开成向量。
```
input_image = torch.rand(64,28,28)
flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())
```
  假设我们输入一个mini-batch为64的28*28的灰度图，我们先将其展开成向量，现在我们就会获得一个mini batch的长度为784的向量。
  
  torch.Size([64, 784])
- 之后就是使用nn.Sequential创建网络结构。我们可以看到是三个线性层夹着两个ReLU激活函数。
  
  nn.Sequential可以看做是模型层的容器，我们在其中按顺序定义模型的每一层。
  - nn.Linear就是最简单的线性层，输入向量，将其和权重以及bias进行计算之后获得输出，计算为$H= WX+b$。
```
layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())
```
    这一步我们是在第一个隐藏层的计算中将特征压缩到20维，所以输出结果为：
    
    torch.Size([64, 20])
  - nn.ReLU就是激活层，将线性运算转化为非线性，可以让神经网络学到更广泛的知识。
    
    在这个模型里我们激活函数用的就是ReLU，当然还有其他的激活函数，常见可以看这里：常用激活函数 - 掘金 (juejin.cn)
```
print(f"Before ReLU: {hidden1}\n\n")
hidden1 = nn.ReLU()(hidden1)
print(f"After ReLU: {hidden1}")
```
    上边这段代码就是输出一下使用激活函数前后的结果。因为batch size 设置的比较大，我这就不展示结果了。
forward：对输入数据的操作都是在这里实现的。
- 先使用self.flatten对输入x进行展开。
- 结果使用logits存储，最后返回结果。

上一步我们已经创建好网络，我们把模型丢到设备上进行计算，并输出其网络结构。

model = NeuralNetwork().to(device)
print(model)

不出意外的话我们可以看到如下的网络结构：

NeuralNetwork( $\quad$(flatten): Flatten(start_dim=1, end_dim=-1) $\quad$(linear_relu_stack): Sequential( $\quad\quad$ (0): Linear(in_features=784, out_features=512, bias=True) $\quad\quad$ (1): ReLU() $\quad\quad$ (2): Linear(in_features=512, out_features=512, bias=True) $\quad\quad$ (3): ReLU() $\quad\quad$ (4): Linear(in_features=512, out_features=10, bias=True) $\quad\quad$) $\quad$)

模型建好了，我们当然是要让模型运行了。所以我们要把数据传给模型，Pytorch通过一些背后的操作，会直接调用模型的forward方法，并不需要我们使用model.forward()手动调用。想知道背后原理的可以看源码：pytorch/module.py at 270111b7b611d174967ed204776985cefca9c144 · pytorch/pytorch (github.com)

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

因为没数据集，所以我们在这里直接rand一个X，不是实际图，就是随机生成一个矩阵。

结果logits是用模型计算获得的。获得的结果应该是在$[-\infty,+\infty]$之间的。

因为是个图像分类模型，分类的最后肯定是softmax计算。nn.Softmax会将模型出来的结果做一个scale，将其限制在$[0,1]$之间，用于表示模型分类的结果并将其存在pred_probab里，这个softmax出来的向量结果加起来是等于1的。

这里我们的模型是 未训练的 ，我们直接拿着做预测了，预测结果存在y_pred中，最后print输出一下预测的结果是什么，这里的结果很大概率是在胡说八道的，因为 本文的重点 就是前边那一段创建网络的部分。之后我们会讲如何训练模型。

模型参数

我们神经网络的很多层都是被参数化的，比如层与层之间会有weights和bias，训练过程中会对这些参数进行优化。nn.Module会对其自动追踪，你可以使用parameters()或者 named_parameters()进行访问。

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

秒客网

零基础用PyTorch搭建自己的神经网络

使用GPU加速训练

创建网络

模型参数

相关文章