Полносвязная нейросеть на Pytorch – базовая структура кода на простом примере

Ниже представлен базовый шаблон кода для создания нейросети.

В примере рассмотрено создание двух полносвязных слоев.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from random import randint


# Создание класса для простой нейросети из двух слоев
class NetGirl(nn.Module):
    def __init__(self, input_dim, num_hidden, output_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, num_hidden)
        self.layer2 = nn.Linear(num_hidden, output_dim)

    def forward(self, x):
        x = self.layer1(x) # Пропускаем входные данные через первый слой
        x = F.tanh(x)      # Полученное скал. произведение + байес пропускаем через функцию активации
        x = self.layer2(x) # Пропускаем результат первого слоя через второй слой
        x = F.tanh(x)      # Пропускаем скал. произведение второго слоя через функцию активации
        return x


# Данные для обучения
x_train = torch.FloatTensor([(-1, -1, -1), (-1, -1, 1), (-1, 1, -1), (-1, 1, 1),
                             (1, -1, -1,), (1, -1, 1), (1, 1, -1), (1, 1, 1)])
y_train = torch.FloatTensor([-1, 1, -1, 1, -1, 1, -1, -1])
total = len(y_train)


# Создадим модель и зададим количества нейронов для двух слоев
model = NetGirl(3, 3, 1)

# Обучаем модель
# Определим оптимизатор и функцию потерь
optimizer = optim.RMSprop(params=model.parameters(), lr=0.01)
loss_func = torch.nn.MSELoss()

# Переводим модель в режим обучения
model.train()

# Цикл обучения
for _ in range(1000):
    k = randint(0, total-1)
    y = model(x_train[k])
    y = y.squeeze()    # Добавление первой оси к y (связано с обработкой батчами)
    loss = loss_func(y, y_train[k])

    # Один шаг алгоритма стохастического градиентного спуска
    optimizer.zero_grad()    # обнуление градиентов перед очередным шагом
    loss.backward()
    optimizer.step()

# Переод модели в режим эксплуатации
model.eval()

# Можно отключать и включать локальные градиенты вручную
# model.requires_grad_(False)
# model.requires_grad_(True)

# Оценка модели
for x, d in zip(x_train, y_train):
    with torch.no_grad():    # При тестировании модели рекомендуется отключать градиенты для экономии памяти
        y = model(x)
        print(f"Выходное значение НС: {y.data} => {d}")