Использование CPU и GPU на примере простой нейронной сети с использованием PyTorch

В этом уроке мы рассмотрим, как можно использовать CPU и GPU для обучения нейронных сетей на примере простой нейронной сети с использованием библиотеки PyTorch. CPU (Central Processing Unit) и GPU (Graphics Processing Unit) – это два основных типа оборудования, которые могут использоваться для выполнения вычислений в машинном обучении.

PyTorch – это популярная библиотека для работы с нейронными сетями, которая обеспечивает высокую производительность и гибкость. Она позволяет легко создавать, обучать и развертывать нейронные сети.

Наша простая нейронная сеть будет состоять из одного скрытого слоя с несколькими нейронами и функцией активации ReLU. Мы будем обучать эту сеть на наборе данных MNIST, который состоит из изображений рукописных цифр.

Шаг 1: Установка PyTorch
Для начала установим PyTorch. Мы можем установить его с помощью pip:

pip install torch torchvision

Шаг 2: Импорт необходимых библиотек
После установки PyTorch импортируем необходимые библиотеки:

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

Шаг 3: Загрузка и предварительная обработка данных MNIST
Загрузим набор данных MNIST и преобразуем его в формат, с которым может работать PyTorch:

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

Шаг 4: Определение модели нейронной сети
Определим простую нейронную сеть с одним скрытым слоем:

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

net = SimpleNN()

Шаг 5: Определение функции потерь и оптимизатора
Определим функцию потерь и оптимизатор для обучения модели:

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Шаг 6: Обучение модели
Теперь мы можем начать обучение модели. Для этого зададим количество эпох и переберем данные в тренировочном наборе:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)

for epoch in range(5):  # запустим обучение на 5 эпох
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 1000 == 999:  # выводим статистику каждые 1000 мини-пакетов
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 1000))
            running_loss = 0.0

Шаг 7: Проверка точности модели
После завершения обучения можно проверить точность модели на тестовом наборе данных:

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

В этом уроке мы рассмотрели, как использовать CPU и GPU для обучения нейронной сети на примере простой нейронной сети с использованием библиотеки PyTorch. Надеюсь, этот урок был полезен для вас!