Hi all — I’m new to deep learning and PyTorch and am playing with some small examples to learn the ropes. I just wanted to check that I’ve implemented something correctly.
I’m looking to use a CNN to detect whether there’s a quasi-diagonal line in an image of a distance matrix (1 = Yes, 0 = No), and if so where (t_start, t_end), like in the image below.
Screen Shot 2020-11-02 at 8.12.28 PM787×559 209 KB
Here is the code I’ve written:
BCE      = torch.nn.BCELoss()
MSE      = torch.nn.MSELoss()

labels   = torch.tensor([
    # Class, t_start, t_end
    [0, 0, 0],
    [1, 0.25, 0.31],
    [0, 0, 0]
])

# Simulate output = model(distance_matrix)
output = torch.tensor([
    [0.1478, 0.1370, 0.2288],
    [0.7052, 0.3207, 0.5404],
    [0.4511, 0.2978, 0.5279]
])

# Binary cross-entropy error on class labels, mean-squared error on time points
# Use labels to zero-out irrelevant rows
loss = BCE(output[:,:1], labels[:,:1]) + MSE(labels[:,:1] * output[:,-2:], labels[:,-2:])

# etc ...
# loss.backward()
# optimizer.step()

I was wondering whether using the labels to zero out the rows of the output I wasn’t interested in (where class = 0) was the right way to go about this. Thanks!