torch.nn 命名空間提供了構建自己的神經網路所需的所有基本組件。在PyTorch中，每個模組都是 nn.Module 的子類別。神經網路本身也是一個模組，由其他模組（層）組成。

首先我們先引入要使用的模組：

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

檢查是否能夠用硬體加速，如GPU或MPS上訓練我們的模型。讓我們檢查一下是否有torch.cuda或torch.backends.mps可用，否則我們將使用CPU。

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

透過繼承 nn.Module 來定義我們的神經網路，並在 __init 方法中初始化神經網路的層：

class NeuralNetwork(nn.Module):
    def __init__(self, img_size):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(img_size, 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

接著創建了一個名為 NeuralNetwork 的實例，在初始化的時候要給定圖片的大小，因此我輸入224*224，然後將它移到指定的設備上，並顯示出它的結構：

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

創建一個 224*224 的 tensor 輸入給剛剛創建的神經網路，確認是否可以正常輸出：

X = torch.rand(1, 224, 224, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

結語

今天介紹了怎麼簡易建立神經網路，明天會每個細節的作用和使用。