Pytorch学习笔记（五）Learn the Basics - Build the Neural Network

这篇博客瞄准的是 pytorch 官方教程中 Learn the Basics 章节的 Build the Neural Network 部分。

• 官网链接：https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html

完整网盘链接: https://pan.baidu.com/s/1L9PVZ-KRDGVER-AJnXOvlQ?pwd=aa2m 提取码: aa2m 

---

Build the Neural Network

torch.nn 提供了构建神经网络所需的模块。PyTorch 中的每个模块都是 nn.Module 的子类。神经网络本体就是一个由其他模块组成的总模块。这种嵌套结构允许轻松构建和管理复杂的架构。

---

Step1. 导入依赖库

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

---

Step2. 获得本地可以使用的硬件

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

---

Step3. 定义神经网络模型

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__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

# 将模型移动到计算设备上
model = NeuralNetwork()
model.to(device)

---

Step4. 使用模型

这里直接给了一个随机数，生成的内容基本没有意义：

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}")