Pytorch-04 搭建神经网络架构工作流

搭建神经网络架构

在pytorch中，神经网络被抽象成由一系列对数据执行特定操作的层或者模块组成，比如下面的Attention实现，每个块都是一个模块或者层。

如果你想快速搭建网络架构，torch.nn这个命名空间提供了所有很多开箱即用的层/模块/算子：

如果你想自定义一个模块也是完全可以的。每个模块都是nn.Module的子类，你只需要继承然后复写即可，这个后面有例子。

这种简洁的架构抽象可以让使用pytorch的人们快速搭建并管理精妙的模型架构。

接下来，我们将搭建一个神经网络来分类FashionMNIST数据集，来过一遍搭建网络的工作流。

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

1. 获取可能的加速设备

为了在 加速器（accelerator） 上训练我们的模型，例如 CUDA、MPS、MTIA 或 XPU，我们将遵循以下逻辑：

如果当前设备有可用的加速器，我们就使用它；否则，我们将使用 CPU。

device = torch.accelerator.current_accelerator().type if  torch.accelerator.is_available() else "cpu"
print(