Patch based cross entropy loss

username: TonyTr
cooked: <p>I am trying to implement a custom loss function for a variant of ViT, where the output is a prediction for each patch from the original image.</p>
<p>The input is of shape [BxCxHxW] and the label for each image is of shape [BxNxRHxRW] where N is the number of classes, RH and RW are reshaped patch size, this is better understood using the following example:</p>
<p>Input shape is [5x1x512x512] with 4 classes. Computing patches (patch size is 32x32) results in [5x256x1024] (ignoring the class token for now). The output from the transformer is reshaped into [5x4x16x16] which is also the final output of the model.</p>
<p>The label is of shape [5x4x16x16] where the second dimension is a probability vector of size 4 representing the probabilities of each class (this is done for all patches).</p>
<p>I am trying to compute the loss between the output and the label. Using nn.CrossEntropyLoss does not work because of the extra dimension. I created (borrowed from a similar project) a custom loss function:</p>
<pre><code class="lang-python">class PatchCrossEntropy(nn.Module):
    def __init__(self):
        super(PatchCrossEntropy, self).__init__()

    def forward(self, x, target):
        loss = torch.sum(-target * F.log_softmax(x, dim=1), dim=1)
        return loss.mean()
my_criterion = PatchCrossEntropy()
x = np.random.RandomState(42).normal(size=(batch_size,num_patches,num_classes))
x = x/x.sum(axis=2,  keepdims=True)
y = np.random.RandomState(41).normal(size=(batch_size,num_patches,num_classes))
y = y/y.sum(axis=2,  keepdims=True)

label = torch.as_tensor(x)
target = torch.as_tensor(y)

loss1 = my_criterion(label,label)
print(loss1)
</code></pre>
<p><strong>tensor(-269.7687, dtype=torch.float64)</strong></p>
<p>I expect that the output to be 0 but its not. I think that the problem is because of which dim I am summing from. I also don’t know if I should be using the softmax in the first place because the last layer of the model is a softmax layer, so I think an extra one inside the loss function is not needed.</p>
<p>I can share more code if it helps illustrate the problem more, but this is part of a project so some parts of the code I might not be able to share.</p>

username: gphilip
cooked: <aside class="quote no-group" data-username="TonyTr" data-post="1" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>I expect that the output to be 0</p>
</blockquote>
</aside>
<p>Could you explain why you expect the output to be zero?</p>

username: TonyTr
cooked: <p>To test the loss function, I use the same label as an input and as a target, the output is supposed to be 0 because both are the same.</p>

username: gphilip
cooked: <aside class="quote no-group" data-username="TonyTr" data-post="3" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>the output is supposed to be 0 because both are the same.</p>
</blockquote>
</aside>
<p>This is the part I don’t understand. <em>Why</em> do you expect the output of the <em>specific</em> loss function that you have defined in the code above to be zero when both its inputs are identical?</p>

username: TonyTr
cooked: <p>Correct me if im wrong, but as far as I know, CE will output a loss of 0 when both of inputs are the same. Since CE can’t take the dimensions of my tensors as defined above, I am implementing a custom CE loss function that accepts said dimensions. Which is why I am expecting a zero.</p>

username: gphilip
cooked: <aside class="quote no-group" data-username="TonyTr" data-post="5" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>CE will output a loss of 0 when both of inputs are the same</p>
</blockquote>
</aside>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss" rel="noopener nofollow ugc">The PyTorch implementation of Cross Entropy loss</a> does not accept two identical inputs: it expects the shapes of the input and target to be different. Is this not the CE that you are talking about?</p>

username: TonyTr
cooked: <p>I am an idiot. The reason I thought that the output is supposed to be 0 is because my calculations on the paper were wrong, I just copied the results. Sorry bout the confusion.</p>
<p>Now I don’t know how am I supposed to test the custom function, my labels are not class labels, they are just probability vectors (basically the labels are precomputed probability maps), therefore I don’t have the option of using the original CE loss.</p>
<p>I can generate new labels that match the input to pytorch’s CE but then I get an error because the labels are still of bigger dimension than it supports.</p>

username: gphilip
cooked: <aside class="quote no-group" data-username="TonyTr" data-post="7" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>Now I don’t know how am I supposed to test the custom function, my labels are not class labels, they are just probability vectors (basically the labels are precomputed probability maps), therefore I don’t have the option of using the original CE loss.</p>
</blockquote>
</aside>
<p>Will something like <a href="https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html#torch.nn.NLLLoss" rel="noopener nofollow ugc">NLLLoss</a> help?</p>

username: TonyTr
cooked: <p>This might be it, two questions:</p>
<ol>
<li>I need to compute the loss on the first dim, as this dimension contains the probability maps after reshaping ,<code>NLLLoss</code> does not give me an option to compute the loss on a specific dim against the target.</li>
<li>How do I test the loss function against toy inputs? I need to be sure that its exactly what I need.</li>
</ol>

username: TonyTr
cooked: <p>Just tested it, I can’t apply it on my labels because of the shape:<br>
<code>only batches of spatial targets supported (3D tensors) but got targets of dimension: 4</code></p>
<p>It is what I want, I just need to figure out how to implement it on the first dimension.</p>

username: gphilip
cooked: <aside class="quote no-group" data-username="TonyTr" data-post="9" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>I need to compute the loss on the first dim,</p>
</blockquote>
</aside>
<p>I am not sure I understand this. If your labels are a <code>[3,4]</code> tensor, what does this translate to? You need to compute the loss on the first row, on the first column, or something else?</p>
<aside class="quote no-group" data-username="TonyTr" data-post="9" data-topic="131179">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e5b9ba/48.png" class="avatar"> TonyTr:</div>
<blockquote>
<p>How do I test the loss function against toy inputs?</p>
</blockquote>
</aside>
<p>You could use random arrays. I did the following to test the loss function from your code above:</p>
<pre><code class="lang-auto">x = torch.rand(3, 4, 5)
torch.sum(-x * F.log_softmax(x, dim=1), dim=1)
</code></pre>

username: TonyTr
cooked: <p>In the case of <code>[3,4]</code> I would need to compute the loss per row (every row is a dense map). In my case, labels are of shape [5x4x16x16] and I would need to compute the loss along the second dimension.</p>

username: gphilip
cooked: <p>Something like:</p>
<pre><code class="lang-auto">x = torch.rand(5, 4, 16, 16)
for i in range(4):
    print("Computing loss for" + str(x[:, i, :, :]))
</code></pre>
<p>?</p>

username: TonyTr
cooked: <p>Yes, I want to do this operation but vectorised.</p>

username: Zhaoyi-Yan
cooked: <p>Please correct me if I’m mistaken. I’m afraid the concept seems a little strange to me. Assume a tiny dog is laying on the ground. The target label is therefore “dog,” but the majority of the image is “background,” so how do you designate the background region? It appears that dog-label can only be applied to all image patches. Or perhaps you have taken steps to address such a situation.</p>

username: TonyTr
cooked: <p>Correct, in the label generation process, I assign a label to every patch for every class (including background). This is possible because of the dataset I am working with which comes with a segmentation mask that is used during the labelling of every region.</p>