Функция torch.nn.GRU в PyTorch

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

torch.nn.GRUЯвляется модулем Gated Recurrent Unit в PyTorch.

GRU — это упрощённая версия LSTM с меньшим числом параметров, более быстрыми вычислениями и схожей производительностью.

Определение функции

torch.nn.GRU(input_size, hidden_size, num_layers=1, bias=True, batch_first=True, dropout=0, bidirectional=False)

Примеры использования

Пример 1: Базовое использование

Пример

import torch
import torch.nn as nn

# GRU: 256-мерный вход, 256-мерное скрытое состояние, 2 слоя
gru = nn.GRU(input_size=256, hidden_size=256, num_layers=2, batch_first=True)

# Вход: batch=4, sequence=10, features=256
x = torch.randn(4, 10, 256)

output, hidden = gru(x)

print("Форма входа:", x.shape)
print("Форма выхода:", output.shape)
print("Форма скрытого состояния:", hidden.shape)

Пример 2: Сравнение с LSTM

Пример

import torch
import torch.nn as nn
import time

# LSTM и GRU с одинаковой конфигурацией
lstm = nn.LSTM(256, 256, 1, batch_first=True)
gru = nn.GRU(256, 256, 1, batch_first=True)

x = torch.randn(32, 100, 256)

# Сравнение производительности
for model, name in [(lstm, "LSTM"), (gru, "GRU")]:
    start = time.time()
    for _ in range(100):
        _ = model(x)
    print(f"{name} Время: {time.time()-start:.3f}s")

Пример 3: Задача классификации

Пример

import torch
import torch.nn as nn

class GRUClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super(GRUClassifier, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.gru = nn.GRU(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, x):
        embedded = self.embedding(x)
        _, hidden = self.gru(embedded)
        # Объединить конечные скрытые состояния двунаправленной
        hidden = torch.cat([hidden[-2], hidden[-1]], dim=1)
        return self.fc(hidden)

model = GRUClassifier(10000, 128, 128, 2)
x = torch.randint(0, 10000, (8, 50))
output = model(x)

print("Вход:", x.shape, "-> Выход:", output.shape)

LSTM против GRU

Аспект LSTM GRU
Количество параметров Подробнее Меньше
Гейтирование 3 вентиля 2 вентиля
Вычисления Медленнее Быстрее

Сценарии использования

  • Моделирование последовательностей: текст, аудио
  • Быстрое прототипирование: когда ресурсы ограничены
  • Машинный перевод: сторона энкодера

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

Другие расширения