Метод retain_grad

Метод retain_grad класса Tensor в PyTorch позволяет сохранить градиент для промежуточного тензора, который не является листом (leaf) в вычислительном графе. По умолчанию PyTorch хранит градиенты только для листовых тензоров (тех, которые были созданы явно и имеют requires_grad=True), чтобы оптимизировать использование памяти. Однако в некоторых случаях, например, при отладке или для анализа градиентов в промежуточных слоях нейронной сети, может потребоваться доступ к градиентам этих тензоров.

Метод вызывается у тензора без аргументов. После вызова retain_grad, градиент для этого тензора будет накапливаться в атрибуте grad во время выполнения обратного распространения.

Синтаксис

tensor.retain_grad()

Пример

Рассмотрим базовый пример использования метода для промежуточного тензора:

import torch

t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
res = t * 2
res.retain_grad()
loss = res.sum()
loss.backward()

print(res.grad)

Результат выполнения кода:

tensor([1., 1., 1.])

В данном примере тензор res является промежуточным, но благодаря вызову retain_grad мы можем получить его градиент после выполнения backward.

Пример

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

import torch

t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
res = t * 2
loss = res.sum()
loss.backward()

print(res.grad)

Результат выполнения кода:

None

Как видно, без вызова retain_grad градиент промежуточного тензора res равен None.

Пример

Метод часто используется для анализа градиентов внутри сложных моделей. Рассмотрим пример с моделью, состоящей из двух линейных слоев:

import torch
import torch.nn as nn

torch.manual_seed(0)

class Model(nn.Module):
	def __init__(self):
		super().__init__()
		self.l1 = nn.Linear(3, 3)
		self.l2 = nn.Linear(3, 1)
	
	def forward(self, x):
		x = self.l1(x)
		x.retain_grad()
		x = self.l2(x)
		return x

model = Model()
inputs = torch.randn(1, 3, requires_grad=True)
outputs = model(inputs)
loss = outputs.sum()
loss.backward()

print(model.l1.output.grad)

Результат выполнения кода:

tensor([[-0.3600, 0.1198, 0.1898]])

В этом примере мы сохранили градиент выхода первого слоя, чтобы проанализировать его в процессе обучения.

Смотрите также

атрибут requires_grad,
который указывает, нуждается ли тензор в вычислении градиента
метод backward,
который вычисляет градиенты для всех тензоров с requires_grad=True
атрибут grad,
в котором хранится вычисленный градиент тензора
атрибут grad_fn,
который указывает на функцию, создавшую тензор в графе вычислений