在PyTorch中搭建神經網絡通常涉及以下步驟:
導入必要的庫:
import torch
import torch.nn as nn
import torch.optim as optim
定義網絡結構:
你可以創建一個繼承自nn.Module
的類來定義你的網絡結構。例如,一個簡單的全連接神經網絡可以如下定義:
class SimpleNN(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, output_size)
def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out
初始化網絡、損失函數和優化器:
input_size = 784 # 假設輸入是一個28x28的圖像
hidden_size = 128
output_size = 10 # 假設輸出是10個類別的概率分布
net = SimpleNN(input_size, hidden_size, output_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)
準備數據集: 你需要將數據集加載到內存中,并進行必要的預處理。例如,對于圖像數據,你可能需要將其展平為一維向量,并進行歸一化。
from torchvision import datasets, transforms
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
訓練網絡:
num_epochs = 10
for epoch in range(num_epochs):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
測試網絡:
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the test images: {100 * correct / total:.2f}%')
以上就是在PyTorch中搭建和訓練一個簡單神經網絡的步驟。你可以根據自己的需求調整網絡結構、損失函數和優化器,以及數據預處理的方式。