Почему модель на обучении занимает больше памяти, чем на инференсе?

Почему модель на обучении занимает больше памяти, чем на инференсе?

Это фундаментальное различие в том, как работает обучение нейронной сети по сравнению с инференсом. Давайте разберём все причины.

Основные причины увеличения потребления памяти при обучении

1. Сохранение активаций для backpropagation

Во время обучения необходимо вычислить градиенты по отношению к каждому параметру модели. Для этого нужны активации (выходы каждого слоя) из фазы forward pass. Эти активации сохраняются в памяти на протяжении всего forward pass, затем используются при обратном распространении ошибки (backpropagation).

import torch
import torch.nn as nn

# Пример: при forward pass сохраняются все промежуточные активации
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1000, 500)
        self.fc2 = nn.Linear(500, 100)
    
    def forward(self, x):
        # Активация fc1 сохраняется для backprop
        x = self.fc1(x)
        x = torch.relu(x)
        # Активация fc2 сохраняется для backprop
        x = self.fc2(x)
        return x

# На инференсе активации не нужны
model.eval()
with torch.no_grad():
    output = model(input_data)  # Активации не сохраняются

2. Градиенты параметров

При обучении нужно хранить градиенты для каждого параметра модели. Если модель имеет N параметров, то нужно дополнительно N памяти для хранения градиентов.

# При обучении
loss.backward()  # Вычисляются градиенты

# Память расходуется на:
# 1. Сами параметры (веса, смещения)
# 2. Градиенты параметров
# 3. Активации для backprop

3. Оптимизаторы с состоянием

Оптимизаторы типа Adam, RMSprop хранят состояние для каждого параметра:

import torch.optim as optim

# Adam хранит для каждого параметра:
# 1. Первый момент (m) — скользящее среднее градиентов
# 2. Второй момент (v) — скользящее среднее квадратов градиентов

optimizer = optim.Adam(model.parameters())
# Память = параметры + градиенты + m + v
# Итого: 4x от размера параметров

# SGD с momentum требует меньше: параметры + градиенты + momentum
optimizer_sgd = optim.SGD(model.parameters(), momentum=0.9)

4. Batch processing

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

# Во время обучения: батч целиком в памяти
batch_size = 32
input_size = 1000
batch = torch.randn(batch_size, input_size)  # 32 * 1000 элементов

# На инференсе часто используется меньший батч или даже один пример
single_sample = torch.randn(1, input_size)

Сравнение потребления памяти

Для модели с N параметрами:

• Инференс (eval mode): примерно N памяти (сами параметры)
• Обучение (training mode):
  • Параметры: N
  • Градиенты: N
  • Активации: ~2-5N (зависит от глубины и ширины сети)
  • Состояние оптимизатора: 2N (Adam) или N (SGD с momentum)
  • Итого: 5-9N памяти

Способы снижения потребления памяти при обучении

# 1. Gradient Checkpointing (торговля памятью на время)
from torch.utils.checkpoint import checkpoint

class ModelWithCheckpointing(nn.Module):
    def forward(self, x):
        # Не сохраняет активации, пересчитывает при backprop
        x = checkpoint(self.layer1, x)
        x = checkpoint(self.layer2, x)
        return x

# 2. Использование меньшего батша
batch_size = 8  # Вместо 32

# 3. Mixed precision training
from torch.cuda.amp import autocast
with autocast():
    output = model(input_data)

# 4. Более простой оптимизатор
optimizer = optim.SGD(model.parameters())  # Вместо Adam

Практический пример

import torch

# Модель: 1M параметров
param_size = 1_000_000
param_memory = param_size * 4 / (1024**3)  # В ГБ (float32)

# Инференс: ~4 ГБ
infer_memory = param_memory

# Обучение с Adam:
train_memory = param_memory * 9  # ~36 ГБ
# - Параметры: 4 ГБ
# - Градиенты: 4 ГБ
# - m (Adam): 4 ГБ
# - v (Adam): 4 ГБ
# - Активации и батч: ~16 ГБ

Вывод: обучение требует в 5-9 раз больше памяти, чем инференс, из-за необходимости хранить активации, градиенты и состояние оптимизатора. Это один из основных вызовов при обучении больших моделей.