Unet Multiclassification : Output of my prediction is not what i expected

username: nestlee
cooked: <p>I am building a unet architecture for multiclass classification. The goal is to recognize food on images (eggs, tomato, cheese, etc). I am using Python and Pytorch as library. I have run my train code with my unet architecture but the output prediction is not what i expected. The output show only one class of value 1. Can someone help me with this ? Thanks a lot</p>
<p>the output :<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/b/2b9154bf39d0cf6a4d59d5bbbf7634951428db59.png" alt="output" data-base62-sha1="6dpTTzeMWpzV9idRFBCvsO4hr85" width="640" height="480"></p>
<p>Here is my unet architecture :</p>
<pre><code class="lang-auto">def double_conv(in_channels, out_channels):
    conv = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
        nn.ReLU(inplace=True)
    )
    return conv


class UNet(nn.Module):
    def __init__(self, num_classes):
        super(UNet, self).__init__()

        # definition of max pooling :

        self.max_pool_2x2_1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.max_pool_2x2_2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.max_pool_2x2_3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.max_pool_2x2_4 = nn.MaxPool2d(kernel_size=2, stride=2)
        # self.max_pool_2x2_5 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.down_conv_1 = double_conv(3, 16)
        self.down_conv_2 = double_conv(16, 32)
        self.down_conv_3 = double_conv(32, 64)
        self.down_conv_4 = double_conv(64, 128)
        # self.down_conv_5 = double_conv(128, 256)
        # self.down_conv_6 = double_conv(256, 512)

        # self.up_conv_trans1 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=2, stride=2)
        # self.up_conv_1 = double_conv(512, 256)

        # self.up_conv_trans2 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=2, stride=2)
        # self.up_conv_2 = double_conv(256, 128)

        self.up_conv_trans3 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=2, stride=2)
        self.up_conv_3 = double_conv(128, 64)

        self.up_conv_trans4 = nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=2, stride=2)
        self.up_conv_4 = double_conv(64, 32)

        self.up_conv_trans5 = nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=2, stride=2)
        self.up_conv_5 = double_conv(32, 16)

        self.out = nn.Conv2d(in_channels=16, out_channels=num_classes,  kernel_size=1)

    def forward(self, image):

        out_conv1 = self.down_conv_1(image)
        out_pool1 = self.max_pool_2x2_1(out_conv1)
        out_conv2 = self.down_conv_2(out_pool1)
        out_pool2 = self.max_pool_2x2_2(out_conv2)
        out_conv3 = self.down_conv_3(out_pool2)
        out_pool3 = self.max_pool_2x2_3(out_conv3)
        out_conv4 = self.down_conv_4(out_pool3)
        # out_pool4 = self.max_pool_2x2_4(out_conv4)
        # out_conv5 = self.down_conv_5(out_pool4)

        # decoder part

        # out_up_conv = self.up_conv_trans2(out_conv5)
        # out_up_conv = self.up_conv_2(torch.cat([out_up_conv, out_conv4], 1))

        out_up_conv = self.up_conv_trans3(out_conv4)
        out_up_conv = self.up_conv_3(torch.cat([out_up_conv, out_conv3], 1))

        out_up_conv = self.up_conv_trans4(out_up_conv)
        out_up_conv = self.up_conv_4(torch.cat([out_up_conv, out_conv2], 1))

        out_up_conv = self.up_conv_trans5(out_up_conv)
        out_up_conv = self.up_conv_5(torch.cat([out_up_conv, out_conv1], 1))

        out_up_conv = self.out(out_up_conv)
        return out_up_conv
</code></pre>
<p>and here is my training code <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<pre><code class="lang-auto">import torch.nn.functional as F
import os
import torch.optim as optim
import numpy as np
from UNET_19classes import *
import random
from data_augmentation import transform


def check_nan(model, input):
    output = model(input)
    return torch.isnan(output).any()

# Sauvegarder le meilleur modèle
def save_model(epoch, model, optimizer, loss):
    print(f"save final model")
    torch.save({
        'epoch': epoch,
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'loss': loss,
    }, f'C:/Users/ne/Desktop/food-recognition/modele_sauvegarder/modele_{epoch}_{loss}.pth')

# Définition des classes
classes = ['background', 'bun', 'laitue', 'tomates', 'cheddar', 'hamburger',
    'sauce', 'oignons', 'cornichon', 'fromage', 'poulet',
    'fish', 'bacon', 'oeuf', 'oignon cuit', 'pomme de terre / frite',
    'avocat', 'crevette', 'poulet pané/tenders', 'champignon']


# Convertir la liste de classes en un objet PyTorch Tensor
class_to_idx = {classes[i]: i for i in range(len(classes))}
idx_to_class = {i: classes[i] for i in range(len(classes))}
class_to_idx_tensor = torch.tensor(list(class_to_idx.values()))


# Définition du modèle UNet
model = UNet(num_classes=20)

# Charger les données d'entraînement et de test
train_source = np.load("images_burger_train_set.npy")
train_source = train_source[..., ::-1]
# train_source = train_source / 255.0
test_source = np.load("images_burger_test_set.npy")
train_target = np.load("mapped_masks.npy")
test_target = np.load("mapped_masks_test.npy")

# Mise aux dimensions correctes [N, C, H, W]
train_target = np.expand_dims(train_target, axis=3)
test_target = np.expand_dims(test_target, axis=3)
arr_zeros = np.zeros((156, 256, 256, 2))
arr_zeros2 = np.zeros((45, 256, 256, 2))
train_target = np.concatenate((train_target, arr_zeros), axis=-1)
test_target = np.concatenate((test_target, arr_zeros2), axis=-1)

# Définition de la fonction perte et de l'optimiseur
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

# Paramètres généraux
num_epochs = 30

correct = 0
total = 0

# Entraînement sur CUDA du modèle
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.cuda()



# Boucle d'entraînement principale
for epoch in range(num_epochs):
    # Phase d'entraînement
    model.train()
    batch_size = 12
    train_acc = 0.0

    # Paramètres d'entraînement
    number_slice_trainset = list(range(156))
    index_train = 0
    number_batch_train = 156 // batch_size
    random.shuffle(number_slice_trainset)
    num_correct_train = 0
    train_loss_array = []
    train_accuracy_array = np.array([0])
    train_loss = 0.0

    # Paramètres de test
    test_epoch_loss = 0
    number_slice_testset = list(range(45))
    index_test = 0
    number_batch_test = 45 // batch_size
    random.shuffle(number_slice_testset)
    num_correct_test = 0
    test_loss_array = []
    test_accuracy_array = np.array([0])
    best_test_loss = np.finfo(np.float32).max
    val_loss = 0.0

    # Boucle d'entraînement sur les lots/batch
    for data in range(number_batch_train):
        # Préparation des données
        batch_train_source_numpy = train_source[number_slice_trainset[index_train:index_train + batch_size], :, :]
        batch_train_source_numpy = np.transpose(batch_train_source_numpy, (0, 3, 1, 2))
        batch_train_source_numpy = batch_train_source_numpy/255.0
        batch_train_target_numpy = train_target[number_slice_trainset[index_train:index_train + batch_size], :, :]
        batch_train_target_numpy = np.swapaxes(batch_train_target_numpy, 1, 3)
        batch_train_target_numpy = np.swapaxes(batch_train_target_numpy, 2, 3)

        batch_train_source_cuda = torch.tensor(batch_train_source_numpy).to(device=device, dtype=torch.float32)
        batch_train_target_cuda = torch.tensor(batch_train_target_numpy).to(device=device, dtype=torch.float32)
        batch_train_target_cuda = batch_train_target_cuda.cuda().float()
        batch_train_target_cuda = torch.argmax(batch_train_target_cuda, dim=1)
        result = transform(batch_train_source_cuda, batch_train_target_cuda)

        batch_train_logits_cuda = model(result[0])

        batch_prediction_cuda = torch.softmax(batch_train_logits_cuda, dim=1)
        preds = batch_prediction_cuda.cpu().detach().numpy()
        # Conversion des probabilités en indice de classe
        preds_idx = np.argmax(preds, axis=1)

        # Remplacement des noms de classes par les indices correspondants
        names_for_each_pixel = preds_idx.astype(np.uint8)

        loss = loss_function(batch_train_logits_cuda, batch_train_target_cuda)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        train_loss += loss.item()
        index_train += batch_size

        preds = F.softmax(batch_train_logits_cuda, dim=1).argmax(dim=1)
        num_correct_train += (preds == batch_train_target_cuda).sum()

    train_loss /= number_batch_train
    print(f"epoch {epoch + 1}, train loss: {train_loss}")
    train_acc = num_correct_train / ((torch.numel(batch_train_logits_cuda)) * number_batch_train)


    # Test du modèle sur les données de validation
    model.eval()
    val_acc = 0.0

    # Boucle de validation sur les lots
    with torch.no_grad():
        for data in range(number_batch_test):
            batch_test_source_numpy = test_source[number_slice_testset[index_test:index_test + batch_size], :, :]
            batch_test_source_numpy = np.swapaxes(batch_test_source_numpy, 1, 3)

            batch_test_target_numpy = test_target[number_slice_testset[index_test:index_test + batch_size], :, :]
            batch_test_target_numpy = np.swapaxes(batch_test_target_numpy, 1, 3)

            batch_test_source_cuda = torch.tensor(batch_test_source_numpy).to(device=device, dtype=torch.float32)
            batch_test_target_cuda = torch.tensor(batch_test_target_numpy).to(device=device, dtype=torch.float32)
            batch_test_target_cuda = torch.argmax(batch_test_target_cuda, dim=1)

            outputs = model(batch_test_source_cuda)

            loss = loss_function(outputs, batch_test_target_cuda)
            test_epoch_loss += loss
            index_test += batch_size

            val_loss += loss.item()
            preds = F.softmax(outputs, dim=1).argmax(dim=1)
            acc = (preds == batch_test_target_cuda).float().mean()
            val_acc += acc.item()

    test_epoch_loss /= number_batch_test
    print(f"epoch {epoch + 1}, test loss: {test_epoch_loss}")
    test_loss_array = np.append(test_loss_array, test_epoch_loss.item())
    np.save(os.path.join(numpy_array_path, 'test_loss_array'), test_loss_array)
</code></pre>

username: KFrank
cooked: <p>Hi Nestlee!</p>
<aside class="quote no-group quote-modified" data-username="nestlee" data-post="1" data-topic="178359" data-full="true">
<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/n/9e8a1a/48.png" class="avatar"> nestlee:</div>
<blockquote>
<p>The output show only one class of value 1.</p>
<pre><code class="lang-auto">        self.up_conv_5 = double_conv(32, 16)

        self.out = nn.Conv2d(in_channels=16, out_channels=num_classes,  kernel_size=1)
</code></pre>
</blockquote>
</aside>
<p>You are making predictions for <code>out_channels = num_classes = 20</code>, but<br>
<code>up_conv_5</code> has already reduced your number of channels to <code>16</code>.  I wouldn’t<br>
expect this to fully explain your issue, but I could see it hurting the performance<br>
of your model in that <code>out</code> is not being passed enough per-pixel information<br>
(only 16 channels) to make full predictions for all 20 classes.  Maybe you could<br>
reduce your numbers of channels more slowly and try to feed something like<br>
48 channels into <code>out</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">from data_augmentation import transform
...
        result = transform(batch_train_source_cuda, batch_train_target_cuda)

        batch_train_logits_cuda = model(result[0])
...
        loss = loss_function(batch_train_logits_cuda, batch_train_target_cuda)
</code></pre>
</blockquote>
</aside>
<p>You don’t show us what <code>transform()</code> actually is, but it looks like it is returning,<br>
perhaps, a tuple, so I speculate that it is transforming your <code>train_source</code> and<br>
<code>train_target</code> together in a consistent way.  You pass the transformed “source”<br>
data to your model to get your predicted “logits.”</p>
<p>But you then pass your <em>untransformed</em> target data to <code>loss_function()</code>.<br>
Depending on what <code>transform()</code> does, your untransformed targets might<br>
no longer be consistent with your “source” and predictions, which could<br>
definitely screw things up.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        preds = F.softmax(batch_train_logits_cuda, dim=1).argmax(dim=1)
</code></pre>
</blockquote>
</aside>
<p>As an aside:  <code>softmax()</code> doesn’t change the <em>order</em> of the values, so using<br>
<code>softmax()</code> won’t affect the result of <code>argmax()</code>.  Therefore, for the purpose<br>
of getting integer class-label predictions from <code>argmax()</code>, you can safely<br>
leave <code>softmax()</code> out.  (Of course, leaving <code>softmax()</code> in won’t change the<br>
result, so it’s okay either way.)</p>
<p>For debugging, I have the following suggestions:</p>
<p>You’ve displayed an “image” of your post-<code>argmax()</code> class labels.  You should<br>
also look at your (pre-<code>argmax()</code>) “logit” “image.”  Are all of per-pixel “logit”<br>
values the same, or do they just map to the same class label?</p>
<p>Visualize some prediction “images” before you’ve trained your model at all<br>
(using the model’s random initialization).  The results should look random,<br>
and certainly shouldn’t predict just a single class.</p>
<p>Also look at some target images.  Do they make sense or do all (or almost all)<br>
of the target pixels have the same value?</p>
<p>Can you overfit a small batch (of say four or eight samples) of source / target<br>
pairs, leaving <code>transform()</code> out?  If you train a lot on the same small batch,<br>
you should be able to train your model to predict (by “memorization”) the<br>
correct target values.</p>
<p>When you train (either for real or in your overfitting test) do your predictions<br>
and model parameters change?</p>
<p>Best.</p>
<p>K. Frank</p>

username: nestlee
cooked: <p>Thank you for your reply ! <a class="mention" href="/u/kfrank">@KFrank</a></p>
<p>I found the problem but cannot see any solution with my knowledge.<br>
The problem come from my mask array, whenever I expand the dimension to have (N, C, W, H) (45, 3, 256, 256), the value of pixels are set to 1, so my mask appear with the correct contour but the pixels have a value of 1.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd7c6171c4c44f49d4d0d2bbe517bdbe8261b931.png" alt="test-target" data-base62-sha1="r2gJ1yBPni1TWoCQFRS6WjHgsmt" width="640" height="480"></p>
<p>Do you have any suggestion to expand the dimensions of my mask array without impacting the pixels value ?</p>

username: nestlee
cooked: <p>I tried to add a 4th dimension without specifiying that it should be (3). So instead of (45, 3, 256, 256), I have (45, 1, 256, 256) and the mask appears at it’s original colors</p>

username: KFrank
cooked: <p>Hi Nestlee!</p>
<aside class="quote no-group" data-username="nestlee" data-post="3" data-topic="178359" data-full="true">
<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/n/9e8a1a/48.png" class="avatar"> nestlee:</div>
<blockquote>
<p>I expand the dimension to have (N, C, W, H) (45, 3, 256, 256)</p>
</blockquote>
</aside>
<p>The short story is that you don’t want to expand your mask dimensions.</p>
<p>Based on my assumptions about your use case, you want<br>
<code>batch_train_target_cuda</code> (the <code>target</code> you pass to your <code>CrossEntropyLoss</code><br>
<code>loss_function</code>) to be a <code>LongTensor</code> with shape <code>[45, 256, 256]</code>.  (This will<br>
be for a batch size of <code>45</code> and <code>H = W = 256</code>.)  Note that this has no channels or<br>
classes dimension.  This tensor’s values should be your ground-truth integer<br>
class labels that run from <code>0</code> to <code>num_classes - 1 = 19</code>.</p>
<p>I speculate that your <code>train_target</code> numpy array already starts out in this form.<br>
If so, all you should have to do is convert it to a pytorch <code>LongTensor</code>.  Do note<br>
that because of how <code>CrossEntropyLoss</code> works, it is important that the <code>target</code><br>
you pass in is a <code>LongTensor</code> and not, for example, a <code>FloatTensor</code> that happens<br>
to consist of integer values.</p>
<p>Best.</p>
<p>K. Frank</p>

username: nestlee
cooked: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a></p>
<p>Why do we have to use LongTensor for this case ? and what about the shape of image tensor (batch_train_source_cuda) ? Does it have to be (N, W, H) too?<br>
At the end, how will we know which class each pixel belongs to if we don’t put in an image tensor with a defined number of channels?<br>
If batch_train_target_cuda is (N, W, H) and batch_train_source_cuda is (N, C, W, H), I think that it will not work.</p>
<p>Thank you,<br>
Kind Regards,</p>

username: KFrank
cooked: <p>Hi Nestlee!</p>
<aside class="quote no-group" data-username="nestlee" data-post="6" data-topic="178359" data-full="true">
<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/n/9e8a1a/48.png" class="avatar"> nestlee:</div>
<blockquote>
<p>Why do we have to use LongTensor for this case ?</p>
</blockquote>
</aside>
<p>Read the <a href="https://pytorch.org/docs/2.0/generated/torch.nn.CrossEntropyLoss.html" rel="noopener nofollow ugc">documentation for CrossEntropyLoss</a> carefully (and see below).</p>
<aside class="quote no-group">
<blockquote>
<p>what about the shape of image tensor (batch_train_source_cuda) ? Does it have to be (N, W, H) too?</p>
</blockquote>
</aside>
<p>No.  For your U-Net implementation (which is typical) the image you input to<br>
your model should have a channels dimension, in your case, three channels.<br>
This is what’s expected by the first layer of your model which is a <code>Conv2</code>.</p>
<p>In your case, it looks like you are using RGB images – three channels – but<br>
even if you were working with single-channel images (e.g., gray-scale), you<br>
would still need a channels dimension, even though it would be a “singleton”<br>
dimension (that is of size <code>1</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>If batch_train_target_cuda is (N, W, H) and batch_train_source_cuda is (N, C, W, H), I think that it will not work.</p>
</blockquote>
</aside>
<p>It will work as follows:</p>
<p>Your <code>CrossEntropyLoss</code> loss function takes an <code>input</code> tensor (the output<br>
of your model – that is the predictions your model makes) of shape<br>
<code>[N, C, W, H]</code> with, in your case, <code>C = num_classes = 20</code>.  It also takes<br>
a <code>target</code> tensor – your labels – with shape <code>[N, W, H]</code>, specifically with<br>
no channels / classes dimension.  But how then does your <code>target</code> tensor<br>
know how many classes you have?  Because its <em>values</em> are integer class<br>
labels in the range <code>[0, num_classes - 1 = 19]</code>.  This will have type <code>long</code>,<br>
that is, will be a <code>LongTensor</code>.</p>
<p>Your model with take as input an image of shape <code>[N, channels = 3, W, H]</code><br>
and output a prediction tensor of shape <code>[N, num_classes = 20, W, H]</code>.<br>
These tensors will have type <code>float</code>, that is, will be <code>FloatTensor</code>s.</p>
<p>It is <code>CrossEntropyLoss</code> that matches up a <code>FloatTensor</code> <code>input</code> of<br>
shape <code>[N, num_classes, W, H]</code> with a <code>LongTensor</code> <code>target</code> of shape<br>
<code>[N, W, H]</code>.*</p>
<p>As an aside, we usually call the last two dimensions of these tensors <code>H</code>,<br>
for height, and <code>W</code>, for width, rather than <code>W</code> and <code>H</code>.  But that doesn’t matter<br>
for the actual computation – it’s just terminology.</p>
<p>*)  <code>CrossEntropyLoss</code> works in two modes:  It can take a <code>target</code> of “hard”<br>
integer class labels of shape <code>[N, W, H]</code> (no <code>num_classes</code> dimension) of<br>
type <code>LongTensor</code> or one of “soft” probabilistic class probabilities of shape<br>
<code>[N, num_classes, W, H]</code> of type <code>FloatTensor</code>.  It decides which mode<br>
to use based on the type and shape of the <code>target</code> tensor you pass into it.<br>
My assumption is that you are using a <code>target</code> that consists of “hard” integer<br>
class labels.</p>
<p>Best.</p>
<p>K. Frank</p>

username: nestlee
cooked: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a> !</p>
<p>Thank you for your explaination !</p>
<p>I modified my code so that the mask appear with the correct classes.</p>
<p>The problem that occurs actually is that the value of test loss gets higher. the parameters I chose are the following :</p>
<ul>
<li>loss function : crossentropy loss</li>
<li>optimizer : adam</li>
<li>lr : 1e-5</li>
<li>batch size : 8 (or 12)</li>
</ul>
<p>why does it increase ? I checked my database (mask image match)</p>

username: KFrank
cooked: <p>Hi Nestlee!</p>
<aside class="quote no-group" data-username="nestlee" data-post="8" data-topic="178359" data-full="true">
<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/n/9e8a1a/48.png" class="avatar"> nestlee:</div>
<blockquote>
<p>The problem that occurs actually is that the value of test loss gets higher.</p>
</blockquote>
</aside>
<p>Let me assume that as you train the loss computed for your training set goes<br>
down (and maybe some performance metric such as accuracy goes up), while<br>
the loss computed for your test set goes up (and maybe some performance<br>
metric goes down).</p>
<p>If so, you are most likely seeing <em>overfitting,</em> where your network gets trained to<br>
“memorize” your training set and give good predictions for it, but it doesn’t really<br>
“learn” to perform well on the data it wasn’t trained on.</p>
<p>Best.</p>
<p>K. Frank</p>