PyTorch训练循环中loss.backward()报错的解决方案咨询

问题背景 #

训练LSTM模型时，调用loss.backward()触发第一个错误：

RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

添加retain_graph=True后又触发第二个错误：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [512, 128]] is at version 2; expected version 1 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

错误原因 #

1. 第一个错误：LSTM的hidden状态被跨batch复用，每次迭代时hidden仍关联上一次的计算图。调用backward()会释放计算图，但下一次迭代还依赖该hidden，导致再次反向传播时计算图已不存在。
2. 第二个错误：添加retain_graph=True后计算图被保留，但LSTM前向传播时会原地修改传入的hidden变量，而保留的计算图仍依赖旧版本的hidden，导致版本不匹配。

解决方法 #

核心是切断hidden与上一次计算图的关联，同时避免原地修改的影响，推荐以下两种方案：

方案一：用detach()分离计算图（推荐） #

每次迭代获取新的hidden后，将其从计算图中分离，这样下一次迭代时hidden不再关联之前的计算图，无需保留计算图。修改训练循环代码：

for epoch in range(initial_num_epochs):
    
    hidden = initial_model.init_hidden(initial_batch_size)

    for i in range(0, initial_training_input_data.size(0) - initial_batch_size, initial_batch_size):
        inputs = initial_training_input_data[i:i + initial_batch_size]
        targets = initial_training_target_data[i:i + initial_batch_size]

        outputs, hidden = initial_model(inputs, hidden)
        # 将隐藏状态和细胞状态都从计算图中分离
        hidden = (hidden[0].detach(), hidden[1].detach())
        
        loss = initial_criterion(outputs, targets)
        
        initial_optimizer.zero_grad()
        # 无需保留计算图，直接反向传播
        loss.backward()
        
        initial_optimizer.step()
        

    print(f'Epoch [{epoch + 1}/{initial_num_epochs}], Loss: {loss.item():.4f}')

方案二：每个batch重新初始化hidden（仅适合无序列记忆需求的场景） #

如果任务不需要LSTM保留跨batch的序列记忆，可以在每个batch前重新初始化hidden，但这会丢失LSTM的时序特性，一般不用于文本序列任务：

for epoch in range(initial_num_epochs):

    for i in range(0, initial_training_input_data.size(0) - initial_batch_size, initial_batch_size):
        # 每个batch前重新初始化hidden
        hidden = initial_model.init_hidden(initial_batch_size)
        inputs = initial_training_input_data[i:i + initial_batch_size]
        targets = initial_training_target_data[i:i + initial_batch_size]

        outputs, hidden = initial_model(inputs, hidden)
        loss = initial_criterion(outputs, targets)
        
        initial_optimizer.zero_grad()
        loss.backward()
        
        initial_optimizer.step()
        

    print(f'Epoch [{epoch + 1}/{initial_num_epochs}], Loss: {loss.item():.4f}')

关键说明 #

• LSTM的hidden包含隐藏状态和细胞状态，两者都需要执行detach()，否则都会关联计算图。
• 不需要使用retain_graph=True，detach()已经切断了循环依赖，每次反向传播后计算图可正常释放。
• 原地修改问题源于PyTorch的LSTM模块会直接更新传入的hidden变量，分离计算图后旧计算图被释放，不再依赖被修改的变量。

内容的提问来源于stack exchange，提问作者mbrain