I have an image segmentation (Learned Random Walker), but I would like to generalize it.
# Init the random walker modules rw = RandomWalker(1000, max_backprop=True) # Load data and init target = init_image(imagem) seeds = init_seeds(torch.zeros(1, 60, 59).long(), target) diffusivities = torch.zeros(1, 2, 60, 59, requires_grad=True) # Init optimizer optimizer = torch.optim.Adam([diffusivities], lr=0.5) loss = torch.nn.NLLLoss() # Main overfit loop for it in range(iterations + 1): optimizer.zero_grad() # Diffusivities must be positive net_output = torch.sigmoid(diffusivities) # Random walker output = rw(net_output, seeds) # Loss and diffusivities update m = torch.nn.LogSoftmax(dim=1) output_log = [m(o) for o in output] # ls = loss(output_log, target) ls = loss(output_log, target) ls.backward() optimizer.step()
target of dimension [1, 1, 60, 59],
output will be of dimension [1, nClass, 60, 59], that is, it will be the
nClass probability matrices for each region.
But for example, with this same dimension, I have the image
target in the range [0, 255] with 3 regions,
NLLLoss does not allow me to do this, because I would have
output with dimension [1, 3, 60, 59] and several pixels greater than 3.
Is there anything you can do to get around this, or some loss function that solves this problem?
Any help is welcome, thanks.