[Python] loss.backward() takes 10 times the time of network.forward()

I want to train my model using own loss function, however I noticed that loss.backward() takes 2 seconds and learning iteration without this function takes only 0.2 seconds.

I tried to optimize the calculation graph but it didnt do much. All of the tensors are using cuda.

Network in question


class half2(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3,20,(11,11))
self.relu = nn.ReLU()
self.maxpool = nn.MaxPool2d((2,2))
self.conv2 = nn.Conv2d(20,40,(5,5))
self.conv3 = nn.Conv2d(40,80,(3,3))
self.fc1 = nn.Linear(80,160)
self.fc2 = nn.Linear(160,80)
def forward(self,img):
res = self.relu(self.conv1(img))
#res = self.maxpool(res)
res = self.relu(self.conv2(res))
res = self.relu(self.conv3(res))
res = res.view((res.shape[1],res.shape[2],res.shape[0]))
res = self.relu(self.fc1(res))
res = torch.relu(self.fc2(res))
return torch.flatten(res)
def train_forward(self,img):
res = self.relu(self.conv1(img))
#res = self.maxpool(res)
res = self.relu(self.conv2(res))
res = self.relu(self.conv3(res))
res = res.view((res.shape[1],res.shape[2],res.shape[0]))
res = self.relu(self.fc1(res))
res = torch.relu(self.fc2(res))
return torch.flatten(res)

My loss function


class ContrastiveLoss(torch.nn.Module):

def __init__(self, margin=1.0):
super(ContrastiveLoss, self).__init__()
self.margin = margin

def forward(self, output1,output2, Y):
D = tfunc.pairwise_distance(output1, output2)
left = Y*0.5
left = left * torch.pow(D,2)

right = (1-Y)*0.5
choose = torch.clamp(self.margin-D, min=0.0)
right = right * torch.pow(choose,2)
return left + right

Learning cycle


running_loss = 0
cnt = 0
same = torch.tensor(0,dtype=torch.float32).to(device)
diff = torch.tensor(1,dtype=torch.float32).to(device)
for epoch in range(5):
for dr1 in range(1,3):
for ig1 in range(1,9):
for dr2 in range(1,3):
for ig2 in range(1,9):
if dr1 == dr2 and ig1 == ig2:
continue
optimizer.zero_grad()
if dr1 == dr2:
target_dist = same
else:
target_dist = diff
img1 = image_to_input(f"./face_dataset/s{dr1}/{ig1}.pgm").to(device)
img2 = image_to_input(f"./face_dataset/s{dr2}/{ig2}.pgm").to(device)

output1 = net3.train_forward(img1)
output2 = net3.train_forward(img2)

loss = criterion(output1,output2,target_dist)
loss.backward()
optimizer.step()
running_loss += loss.item()

cnt += 1

if cnt != 0:
print("epoch",epoch,"set",dr1,"loss",running_loss / cnt)
running_loss = 0.0
cnt = 0

Continue reading...