Mô-đun torch.nn cung cấp công cụ xây dựng mạng nơ-ron. Dựa trên cơ chế autograd, gói này hỗ trợ định nghĩa kiến trúc mạng và tính toán gradient tự động. Lớp nn.Module đóng vai trò nền tảng, chứa các tầng mạng và phương thức forward() xử lý dữ liệu đầu vào.

Quy trình huấn luyện điển hình:

1. Định nghĩa mạng với các tham số học được
2. Lặp qua tập dữ liệu
3. Truyền dữ liệu qua mạng
4. Tính toán hàm mất mát
5. Lan truyền ngược gradient
6. Cập nhật trọng số: trọng_số = trọng_số - tốc_độ_học * gradient

Định nghĩa kiến trúc mạng

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layer1 = nn.Conv2d(1, 6, kernel_size=5)
        self.conv_layer2 = nn.Conv2d(6, 16, kernel_size=5)
        self.linear1 = nn.Linear(16 * 5 * 5, 120)
        self.linear2 = nn.Linear(120, 84)
        self.linear3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.relu(self.conv_layer1(x))
        x = F.max_pool2d(x, kernel_size=2)
        x = F.relu(self.conv_layer2(x))
        x = F.max_pool2d(x, kernel_size=2)
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        return self.linear3(x)

model = ConvNet()
print(model)

Kết quả:

ConvNet(
  (conv_layer1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
  (conv_layer2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (linear1): Linear(in_features=400, out_features=120, bias=True)
  (linear2): Linear(in_features=120, out_features=84, bias=True)
  (linear3): Linear(in_features=84, out_features=10, bias=True)
)

Phương thức backward được tự động sinh bởi autograd. Truy xuất tham số:

params = list(model.parameters())
print(f"Số lượng tham số: {len(params)}")
print(f"Kích thước kernel đầu: {params[0].size()}")

Tính toán truyền thuận

input_sample = torch.randn(1, 1, 32, 32)
output = model(input_sample)
print(output)

Hàm mất mát

target = torch.randn(10).view(1, -1)
loss_fn = nn.MSELoss()
loss = loss_fn(output, target)
print(loss)

Lan truyền ngược

model.zero_grad()
loss.backward()

Cập nhật trọng số

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Vòng lặp huấn luyện
optimizer.zero_grad()
prediction = model(input_sample)
loss_value = loss_fn(prediction, target)
loss_value.backward()
optimizer.step()