Edge Loss function implementation

I am trying to define a loss function to compute the loss between edge reconstruction. The following is my implementation however I suspect I have made some error. I am calculating the edges using convolutions and then performing mse over it.

def edge_loss(out, target, cuda=True):
	x_filter = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
	y_filter = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
	convx = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1, bias=False)
	convy = nn.Conv2d(1, 1, kernel_size=3 , stride=1, padding=1, bias=False)
	weights_x = torch.from_numpy(x_filter).float().unsqueeze(0).unsqueeze(0)
	weights_y = torch.from_numpy(y_filter).float().unsqueeze(0).unsqueeze(0)

	if cuda:
		weights_x = weights_x.cuda()
		weights_y = weights_y.cuda()

	convx.weight = nn.Parameter(weights_x)
	convy.weight = nn.Parameter(weights_y)

	g1_x = convx(out)
	g2_x = convx(target)
	g1_y = convy(out)
	g2_y = convy(target)

	g_1 = torch.sqrt(torch.pow(g1_x, 2) + torch.pow(g1_y, 2))
	g_2 = torch.sqrt(torch.pow(g2_x, 2) + torch.pow(g2_y, 2))

	return torch.mean((g_1 - g_2).pow(2))

The loss is then calculated as follows

loss = edge_loss(out, x)
loss.backward()

I do not want to update the weights of the convolution filters since these are the edge filters needed. Is this implementation correct?

1 Like

If you don’t want trainable features, the functional api is a much better match, so torch.nn.functional.conv2d etc.

Best regards

Thomas

Thanks. But if I use nn.Conv is my method correct ?

@aradhyamathur: Interesting! Could you share your experiment?

you directly use kornia.sobel or kornia.spatial_gradient