持续创作，加速成长！这是我参与「掘金日新计划 · 10 月更文挑战」的第3天，点击查看活动详情
前言
模型是用于处理输入以生成输出的神经网络层的集合，我们可以使用 torch.nn 包来定义模型，nn 包是提供通用深度学习层的模块集合。nn 的一个模块或层接收输入张量，计算输出张量，并保存权重。我们可以使用两种方法在 PyTorch 中定义模型：nn.Sequential 和 nn.Module。在本节中，我们将学习如何定义一个线性层、一个两层网络和一个多层卷积网络。
定义线性层
创建一个线性层并打印出它的输出大小：
from torch import nn
input_tensor = torch.randn(64, 1000)
linear_layer = nn.Linear(1000, 100)
output = linear_layer(input_tensor)
print(output.size())
打印出的线性层输出大小如下：
torch.Size([64, 100])
使用 nn.Sequential 定义模型
我们可以使用 nn.Sequential 通过按顺序传递层来创建深度学习模型。例如，下图中展示的两层神经网络：
如上图所示，网络有四个节点作为输入，隐藏层中包含五个节点，输出层含有一个节点。接下来，我们将介绍如何实现此双层神经网络。
1. 使用 nn.Sequential 实现上图展示的双层神经网络并打印模型：
from torch import nn
model = nn.Sequential(
nn.Linear(4, 5),
nn.ReLU(),
nn.Linear(5, 1),
)
print(model)
上述代码的输出如下：
Sequential(
(0): Linear(in_features=4, out_features=5, bias=True)
(1): ReLU()
(2): Linear(in_features=5, out_features=1, bias=True)
)
使用 nn.Module 定义模型
在 PyTorch 中定义模型的另一种方法是使用 nn.Module 类，在类的 __init__ 方法中指定神经网络层。然后，在 forward 方法中，将神经网络层应用于输入。这种方法为构建定制模型提供了更高的灵活性。以如下多层神经网络模型为例：
如上图所示，该模型有两个卷积层和两个全连接层。接下来，我们学习如何使用 nn.Module 实现该模型。
1. 首先，实现类(继承自 nn.Module )的主体：
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
def forward(self, x):
pass
2. 然后，定义 __init__ 函数：
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4*4*50, 500)
self.fc2 = nn.Linear(500, 10)
3. 接下来，定义 forward 函数：
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4*4*50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
4. 然后，覆盖类函数 __init__ 和 forward：
Net.__init__ = __init__
Net.forward = forward
5. 接下来，创建一个 Net 类的对象并打印模型：
model = Net()
print(model)
以上代码的输出如下：
Net(
(conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
(conv2): Conv2d(20, 50, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear(in_features=800, out_features=500, bias=True)
(fc2): Linear(in_features=500, out_features=10, bias=True)
)
将模型移动到 CUDA 设备
模型是参数的集合，默认情况下，模型在 CPU 中运行。
1. 获取模型的设备：
print(next(model.parameters()).device)
以上代码打印结果如下：
cpu
2. 然后，模型移动到 CUDA 设备：
device = torch.device("cuda:0")
model.to(device)
print(next(model.parameters()).device)
