PyTorchJuniorTechnical

В чём разница между nn.Module и nn.functional?

nn.Module — класс с состоянием (хранит веса, реагирует на train/eval). nn.functional — набор чистых функций без параметров. Для слоёв с весами и Dropout/BatchNorm используйте Module; для stateless операций (relu, softmax) — F.*.

nn.Module и nn.functional: в чём разница

torch.nn.Module — базовый класс для всех слоёв и моделей PyTorch. Он хранит обучаемые параметры (weight, bias), управляет режимами train/eval и интегрируется с оптимизатором. torch.nn.functional — набор чистых функций без состояния: они принимают тензоры и параметры явно и ничего не хранят.

Ключевые отличия

  • Состояние: nn.Linear хранит self.weight и self.bias. F.linear(input, weight, bias) — вы передаёте их сами.
  • train/eval: nn.Dropout и nn.BatchNorm2d автоматически меняют поведение при model.eval(). F.dropout требует явного параметра training=....
  • Регистрация: параметры Module видны в model.parameters() и сохраняются в state_dict. Параметры, переданные в F.*, нужно регистрировать вручную через nn.Parameter.

Пример: одно и то же двумя способами

import torch
import torch.nn as nn
import torch.nn.functional as F

# --- Способ 1: nn.Module ---
class NetModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)   # хранит weight + bias
        self.drop = nn.Dropout(0.3)       # помнит p, реагирует на eval()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))           # relu — stateless, F.* нормально
        x = self.drop(x)
        return self.fc2(x)

# --- Способ 2: nn.functional напрямую ---
class NetFunctional(nn.Module):
    def __init__(self):
        super().__init__()
        # Параметры нужно регистрировать вручную!
        self.w1 = nn.Parameter(torch.randn(256, 784) * 0.01)
        self.b1 = nn.Parameter(torch.zeros(256))
        self.w2 = nn.Parameter(torch.randn(10, 256) * 0.01)
        self.b2 = nn.Parameter(torch.zeros(10))
        self.drop_p = 0.3

    def forward(self, x):
        x = F.relu(F.linear(x, self.w1, self.b1))
        # training=self.training — иначе dropout не выключится при eval()
        x = F.dropout(x, p=self.drop_p, training=self.training)
        return F.linear(x, self.w2, self.b2)

model = NetModule()
model.eval()   # автоматически выключает Dropout и BatchNorm
print(list(model.parameters()))  # видит все 4 тензора

Когда использовать F.*

  • Для операций без состояния: F.relu, F.softmax, F.cross_entropy, F.interpolate.
  • В кастомных слоях, где вы сами контролируете параметры через nn.Parameter.
  • В loss-функциях — F.cross_entropy, F.mse_loss и т.д.

Когда использовать nn.Module

  • Для слоёв с обучаемыми весами: Linear, Conv2d, Embedding.
  • Для слоёв с режимом train/eval: BatchNorm, Dropout, LayerNorm.
  • Когда нужно сохранить/загрузить state_dict.

Подводные камни

  • Использовать F.dropout(...) без training=self.training — dropout не выключится при model.eval(), и предсказания будут случайными.
  • Хранить тензор параметра как обычный атрибут (self.w = torch.randn(...)) вместо nn.Parameter — оптимизатор его не увидит и не обновит.
  • Путать F.batch_norm с nn.BatchNorm2d: у функциональной версии нет running_mean/running_var, которые нужны при инференсе.
  • Вызывать F.softmax перед F.cross_entropycross_entropy уже включает log_softmax внутри, двойное применение даёт неверный лосс.
  • Забыть super().__init__() в nn.Module — все методы (parameters(), to(device)) сломаются с неочевидными ошибками.

Common mistakes

  • Объяснять module vs functional только синтаксисом без shape, dtype, состояния или режима выполнения.
  • Игнорировать leakage, воспроизводимость, пустые входы и скрытые копии данных.
  • Не проверять production-симптомы: latency, память, ретраи, дрейф качества и несовпадение версий.

What the interviewer is testing

  • Может ли связать module vs functional с реальным контрактом входов и выходов.
  • Упоминает ли тесты, метрики, reproducibility и диагностику ошибок.
  • Видит ли различие между demo-кодом в ноутбуке и production-пайплайном.

Sources

Related topics