PyTorch 核心概念与踩坑记录
Quickstart: Creating Models 使用 accelerator 进行加速： 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 import torch import torch.nn as nn device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu" print(f"Using {device} device") # Define model 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) # model 的实际输入是在这里发生 logits = self.linear_relu_stack(x) return logits model = NeuralNetwork().to(device) print(model) 在 PyTorch（以及大多数深度学习框架）的习惯里，forward 函数接收的输入默认是整个 batch 的数据，而不是单个样本。 ...