I’ve been trying to implement a version of an hourglass model but during training I keep running into this error:
RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.
I’ve read the other posts, but unless I’m mistaken none of them solve my problem.
Beyond a custom loss function , I don’t see any reason why it would behave like this since as far as I can tell I’m doing a regular training loop. Although I might have missed something since I’m pretty new to
I’ve been debugging for a couple days now and I haven’t been able make any headway, so someone here might be able to point me in the right direction?
A lot of the training and supplementary code including the loss function i got from here:
https://github.com/bearpaw/pytorch-pose
Training Code:
model.train()
for epoch in range(epochs):
with tch.no_grad():
lr = adjust_learning_rate(optimizer, epoch, lr, [75, 100, 150], 0.34199518933)
for i, (inputload, targetload, meta) in enumerate(training_loader):
input, target = inputload.to(device), targetload.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
with tch.autograd.detect_anomaly():
output = model(input)
print("target", target.size())
print("output", output.size())
loss = criterion(output, target, target_weight)
print("MSE",loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
input.grad = None
target.grad = None
target_weight.grad = None
output.grad = None
target_weight
Loss function:
# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# Modified by Wei Yang (platero.yang@gmail.com)
# ------------------------------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import torch.nn as nn
class JointsMSELoss(nn.Module):
def __init__(self, use_target_weight=True):
super(JointsMSELoss, self).__init__()
self.criterion = nn.MSELoss(reduction='mean')
self.use_target_weight = use_target_weight
def forward(self, output, target, target_weight):
batch_size = output.size(0)
num_joints = output.size(1)
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
loss = 0
for idx in range(num_joints):
heatmap_pred = heatmaps_pred[idx].squeeze()
heatmap_gt = heatmaps_gt[idx].squeeze()
if self.use_target_weight:
loss += 0.5 * self.criterion(
heatmap_pred.mul(target_weight[:, idx]),
heatmap_gt.mul(target_weight[:, idx])
)
else:
loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
return loss / num_joints
Model Code:
import numpy as np
import pandas as pd
import torch as tch
import torch.nn as nn
import pickle
class ConvBlock(nn.Module):
def __init__(self, n, m):
super(ConvBlock, self).__init__()
self.relu = nn.ReLU()
self.bn = nn.BatchNorm2d(n)
print("conv", n, m)
self.conv = nn.Conv2d(n, m, 3, padding=1)
def forward(self, x):
y=self.conv(x)
return y
class Gate(nn.Module):
def __init__(self, n):
super(Gate, self).__init__()
self.module = nn.Linear(n, n, bias=False)
def forward(self, x):
x = x.permute(0,2,3,1)
x = self.module(x)
x = x.permute(0,3,1,2)
return x
class Residual(nn.Module):
def __init__(self,n):
super(Residual, self).__init__()
self.conv1 = ConvBlock(n,n//2)
self.conv2 = ConvBlock(n//2,n//4)
self.conv3 = ConvBlock(n//4,n//4)
self.gate = Gate(n)
def forward(self,x):
y1 = self.conv1(x)
y2 = self.conv2(y1)
y3 = self.conv3(y2)
print("res", y1.size(), y2.size(), y3.size())
y_c = tch.cat((y1,y2,y3),1)
y = tch.add(y_c, self.gate(x))
return y
class ResUp(nn.Module):
def __init__(self,n):
super(ResUp, self).__init__()
self.res = Residual(n)
def forward(self, x, target):
self.upsample = nn.Upsample(target, mode='nearest')
return self.res(self.upsample(x))
class Merge(nn.Module):
def __init__(self,n):
super(Merge, self).__init__()
self.conv = nn.Conv2d(2*n, n, 3, padding=1)
def forward(self, x, y):
print("merge", x.size(), y.size())
return self.conv(tch.cat((x,y),1))
class Hourglass(nn.Module):
def __init__(self,n):
super(Hourglass, self).__init__()
self.resdowns = nn.ModuleList([Residual(n) for i in range(4)])
self.pool = nn.MaxPool2d(2)
self.resups = nn.ModuleList([ResUp(n) for i in range(4)])
self.ress = nn.ModuleList([Residual(n) for i in range(4)])
self.merges = nn.ModuleList([Merge(n) for i in range(4)])
self.skips = []
def forward(self, x):
print("input size", x.size())
for i in range(4):
x = self.resdowns[i](x)
self.skips.append(self.ress[i](x))
x = self.pool(x)
print("skip in", self.skips[i].size(), i)
for i in range(4):
print("skip out", i)
x = self.merges[i](x,self.pool(self.skips[3-i]))
x = self.resups[i](x, self.skips[3-i].size(2))
print("resup", x.size())
print("output size", x.size())
return x
class Model(nn.Module):
def __init__(self,in_channels, hidden_channels, numjoints, num_hg):
super(Model, self).__init__()
self.num_hg = num_hg
self.conv = nn.Conv2d(in_channels, numjoints, 1)
self.modulelist = nn.ModuleList([nn.Sequential(
nn.Conv2d(numjoints, hidden_channels, 1),
Hourglass(hidden_channels),
nn.Conv2d(hidden_channels, numjoints, 1),
) for i in range(num_hg)])
def forward(self, x):
x = self.conv(x)
for i in range(self.num_hg):
x=self.modulelist[i](x)
return x