1 获取训练设备

检查GPU或CPU的可用性，优先使用GPU加速。

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'使用{device}设备')

输出：使用cuda设备

2 定义模型类

通过继承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(784, 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

创建模型实例并移动到设备上。

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

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

3 模型结构解析

通过样本数据了解网络结构。

input_image = torch.rand(3, 28, 28)
print(input_image.size())

输出： torch.Size([3, 28, 28])

4 层功能解析

4.1 Flatten层

将28x28图像展平成784维度数组。

flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())

输出： torch.Size([3, 784])

4.2 Linear层

应用线性变换，映射784维到20维。

layer1 = nn.Linear(784, 20)
hidden1 = layer1(flat_image)
print(hidden1.size())

输出： torch.Size([3, 20])

4.3 ReLU激活

引入非线性，激活层后的输出。

print("Before ReLU:")
print(hidden1)
hidden1 = nn.ReLU()(hidden1)
print("After ReLU:")
print(hidden1)

输出： Before ReLU: tensor([[-0.2541, -0.1397, 0.2342, 0.1364, -0.0437, 0.3759, 0.2808, -0.0619, 0.2780, 0.2830, -0.4725, 0.4298, 0.2717, -0.1618, -0.0604, 0.3242, -0.5874, -0.5922, -0.2481, -0.4181], [-0.1339, -0.1163, 0.1688, 0.1112, 0.1179, 0.3560, 0.0990, -0.1398, 0.2619, -0.1023, -0.7150, -0.1186, 0.3338, -0.0817, 0.1983, -0.2084, -0.3889, -0.2361, -0.0752, -0.2144], [-0.1284, 0.0683, 0.0707, 0.0997, -0.2274, 0.4379, 0.1461, 0.0949, 0.2710, -0.0563, -0.6621, -0.3552, 0.4966, 0.2304, 0.0020, -0.0470, -0.6260, -0.2077, -0.0790, -0.4635])

After ReLU: tensor([[0.0000, 0.0000, 0.2342, 0.1364, 0.0000, 0.3759, 0.2808, 0.0000, 0.2780, 0.2830, 0.0000, 0.4298, 0.2717, 0.0000, 0.0000, 0.3242, 0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.1688, 0.1112, 0.1179, 0.3560, 0.0990, 0.0000, 0.2619, 0.0000, 0.0000, 0.0000, 0.3338, 0.0000, 0.1983, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.0000, 0.0683, 0.0707, 0.0997, 0.0000, 0.4379, 0.1461, 0.0949, 0.2710, 0.0000, 0.0000, 0.0000, 0.4966, 0.2304, 0.0020, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000])

4.4 Sequential容器

通过顺序容器快速搭建网络。

seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(3, 28, 28)
logits = seq_modules(input_image)

4.5 Softmax层

将logits转换为预测概率。

softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)
y_pred = pred_probab.argmax(1)
print(f"预测类别：{y_pred}")

输出：预测类别：tensor([1], device='cuda:0')

5 模型参数

查看模型中的权重和偏置参数。

print("模型结构：", model, "\n\n")
for name, param in model.named_parameters():
    print(f"层：{name} | 参数大小：{param.size()} | 参数值：{param[:2]} \n")

输出：模型结构： NeuralNetwork( (flatten): Flatten(start_dim=1, end_dim=-1) (linear_relu_stack): Sequential( (0): Linear(in_features=784, out_features=512, bias=True) (1): ReLU() (2): Linear(in_features=512, out_features=512, bias=True) (3): ReLU() (4): Linear(in_features=512, out_features=10, bias=True) ) )

层：linear_relu_stack.0.weight | 参数大小：torch.Size([512, 784]) | 参数值：tensor([[-0.0169, 0.0327, -0.0128, ..., -0.0273, 0.0193, -0.0197], [ 0.0309, 0.0003, -0.0232, ..., 0.0284, -0.0163, 0.0171]], device='cuda:0', grad_fn=

层：linear_relu_stack.0.bias | 参数大小：torch.Size([512]) | 参数值：tensor([-0.0060, -0.0333], device='cuda:0', grad_fn=

层：linear_relu_stack.2.weight | 参数大小：torch.Size([512, 512]) | 参数值：tensor([[-0.0294, 0.0120, -0.0287, ..., -0.0280, -0.0299, 0.0083], [ 0.0260, -0.0075, 0.0430, ..., -0.0196, -0.0200, 0.0145]], device='cuda:0', grad_fn=

层：linear_relu_stack.2.bias | 参数大小：torch.Size([512]) | 参数值：tensor([-0.0003, -0.0043], device='cuda:0', grad_fn=

层：linear_relu_stack.4.weight | 参数大小：torch.Size([10, 512]) | 参数值：tensor([[-0.0287, -0.0199, -0.0147, ..., 0.0074, 0.0403, 0.0068], [ 0.0375, -0.0005, 0.0372, ..., -0.0426, -0.0094, -0.0081]], device='cuda:0', grad_fn=

层：linear_relu_stack.4.bias | 参数大小：torch.Size([10]) | 参数值：tensor([-0.0347, 0.0438], device='cuda:0', grad_fn=

6 参考文献

本文内容来源于PyTorch官方文档，具体链接请见原文。

转载地址：http://orxfk.baihongyu.com/

你可能感兴趣的文章