Вы фактически не накапливаете градиенты. Простое отключение optimizer.zero_grad() не имеет никакого эффекта, если у вас есть единственный вызов .backward(), поскольку градиенты изначально равны нулю (технически None, но они будут автоматически инициализированы равными нулю).
Единственная разница между вашими двумя версиями - это то, как вы рассчитываете окончательный проигрыш. For l oop из второго примера выполняет те же вычисления, что и PyTorch в первом примере, но вы делаете их индивидуально, и PyTorch не может оптимизировать (распараллеливать и векторизовать) ваши для l oop, что имеет особенно ошеломляющую разницу. на графических процессорах, при условии, что тензоры не крошечные.
Прежде чем перейти к накоплению градиента, давайте начнем с вашего вопроса:
Наконец, к моему вопросу: что именно происходит 'под hood '?
Каждая операция над тензорами отслеживается в вычислительном графе тогда и только тогда, когда один из операндов уже является частью вычислительного графа. Когда вы устанавливаете requires_grad=True тензора, он создает вычислительный граф с единственной вершиной, сам тензор, который останется листом в графе. Любая операция с этим тензором создаст новую вершину, которая является результатом операции, следовательно, есть ребро от операндов к ней, отслеживающее выполненную операцию.
a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(4.0)
c = a + b # => tensor(6., grad_fn=<AddBackward0>)
a.requires_grad # => True
a.is_leaf # => True
b.requires_grad # => False
b.is_leaf # => True
c.requires_grad # => True
c.is_leaf # => False
Каждый промежуточный тензор автоматически требует градиентов и имеет grad_fn, который является функцией для вычисления частных производных по отношению к его входным данным. Благодаря правилу цепочки мы можем пройти весь граф в обратном порядке, чтобы вычислить производные по каждому отдельному листу, которые являются параметрами, которые мы хотим оптимизировать. Это идея обратного распространения ошибки, также известной как дифференциация обратного режима . Для получения дополнительной информации я рекомендую прочитать Calculus on Computational Graph: Backpropagation .
PyTorch использует именно эту идею, когда вы вызываете loss.backward(), он проходит график в обратном порядке, начиная с loss , и вычисляет производные для каждой вершины. Когда достигается лист, вычисленная производная для этого тензора сохраняется в его атрибуте .grad.
В вашем первом примере это привело бы к:
MeanBackward -> PowBackward -> SubBackward -> MulBackward`
Второй пример: почти идентичны, за исключением того, что вы вычисляете среднее значение вручную, и вместо одного пути для потерь у вас есть несколько путей для каждого элемента расчета потерь. Чтобы прояснить, единый путь также вычисляет производные каждого элемента, но внутренне, что снова открывает возможности для некоторых оптимизаций.
# Example 1
loss = (y - y_hat) ** 2
# => tensor([16., 4.], grad_fn=<PowBackward0>)
# Example 2
loss = []
for k in range(len(y)):
y_hat = model2(x[k])
loss.append((y[k] - y_hat) ** 2)
loss
# => [tensor([16.], grad_fn=<PowBackward0>), tensor([4.], grad_fn=<PowBackward0>)]
В любом случае создается один график с обратным распространением ровно один раз, вот почему это не считается накоплением градиента.
Накопление градиента
Накопление градиента относится к ситуации, когда перед обновлением параметров выполняется несколько обратных проходов. Цель состоит в том, чтобы иметь одни и те же параметры модели для нескольких входных данных (пакетов), а затем обновлять параметры модели на основе всех этих пакетов, вместо того, чтобы выполнять обновление после каждого отдельного пакета.
Давайте вернемся к вашему примеру. x имеет размер [2] , это размер всего нашего набора данных. По какой-то причине нам нужно рассчитать градиенты на основе всего набора данных. Это, естественно, имеет место при использовании пакета размером 2, поскольку у нас будет сразу весь набор данных. Но что произойдет, если у нас будут партии только размером 1? Мы могли бы запускать их по отдельности и обновлять модель после каждого пакета, как обычно, но тогда мы не вычисляем градиенты по всему набору данных.
Что нам нужно сделать, это запустить каждый образец индивидуально с одинаковыми параметрами модели и рассчитать градиенты без обновления модели. Теперь вы можете подумать, разве это не то, что вы сделали во второй версии? Почти, но не совсем, и в вашей версии есть критическая проблема, а именно, что вы используете тот же объем памяти, что и в первой версии, потому что у вас те же вычисления и, следовательно, такое же количество значений в вычислительном графе.
Как освободить память? Нам нужно избавиться от тензоров предыдущего пакета, а также от вычислительного графа, потому что он использует много памяти для отслеживания всего, что необходимо для обратного распространения. Вычислительный граф автоматически уничтожается при вызове .backward() (если не указано retain_graph=True).
def calculate_loss(x: torch.Tensor) -> torch.Tensor:
y = 2 * x
y_hat = model(x)
loss = (y - y_hat) ** 2
return loss.mean()
# With mulitple batches of size 1
batches = [torch.tensor([4.0]), torch.tensor([2.0])]
optimizer.zero_grad()
for i, batch in enumerate(batches):
# The loss needs to be scaled, because the mean should be taken across the whole
# dataset, which requires the loss to be divided by the number of batches.
loss = calculate_loss(batch) / len(batches)
loss.backward()
print(f"Batch size 1 (batch {i}) - grad: {model.weight.grad}")
print(f"Batch size 1 (batch {i}) - weight: {model.weight}")
# Updating the model only after all batches
optimizer.step()
print(f"Batch size 1 (final) - grad: {model.weight.grad}")
print(f"Batch size 1 (final) - weight: {model.weight}")
Вывод (я удалил параметр , содержащий сообщения для удобства чтения):
Batch size 1 (batch 0) - grad: tensor([-16.])
Batch size 1 (batch 0) - weight: tensor([1.], requires_grad=True)
Batch size 1 (batch 1) - grad: tensor([-20.])
Batch size 1 (batch 1) - weight: tensor([1.], requires_grad=True)
Batch size 1 (final) - grad: tensor([-20.])
Batch size 1 (final) - weight: tensor([1.2000], requires_grad=True)
Как видите, модель сохранила один и тот же параметр для всех пакетов, в то время как градиенты накапливались, а в конце есть одно обновление. Обратите внимание, что потери необходимо масштабировать для каждого пакета, чтобы иметь такое же значение для всего набора данных, как если бы вы использовали один пакет.
Хотя в этом примере весь набор данных используется перед выполнением обновления , вы можете легко изменить это, чтобы обновить параметры после определенного количества пакетов, но вы должны не забыть обнулить градиенты после выполнения шага оптимизатора. Общий рецепт:
accumulation_steps = 10
for i, batch in enumerate(batches):
# Scale the loss to the mean of the accumulated batch size
loss = calculate_loss(batch) / accumulation_steps
loss.backward()
if (i - 1) % accumulation_steps == 0:
optimizer.step()
# Reset gradients, for the next accumulated batches
optimizer.zero_grad()
Вы можете найти этот рецепт и другие методы для работы с большими партиями в HuggingFace - Обучение нейронных сетей на больших партиях: Практические советы для 1-GPU, Multi -GPU и распределенные настройки .