pytorch怎么搭建神经网络_企业动态_资讯中心

发布时间:2026-04-25 19:23:11

阅读量:0

在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}%')