How to create a Unet model with target has 3 channels (RGB)

Bruce_Kieu: <p>Hi folks,<br>
I’ve just done the unet model for multi-classes (5) segmentation with targets are gray images (1 channel). I would like to develop the model to targets with 3 channels (RGB), but not success yet. Here are some snippet of the 1 channel targets:</p>
<pre><code class="lang-auto">class myDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        self.images = os.listdir(image_dir) 

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        img_path = os.path.join(self.image_dir, self.images[index])
        mask_path = os.path.join(self.mask_dir, self.images[index])
        image = np.array(Image.open(img_path).convert("RGB"))       
        mask = np.array(Image.open(mask_path).convert("L"), dtype=np.float32)
        # mask = np.array(Image.open(mask_path).convert("RGB"))

        if self.transform is not None:
            augmentations = self.transform(image=image, mask=mask)
            image = augmentations["image"]
            mask = augmentations["mask"]
        
        return image, mask
</code></pre>
<p>What transform function does is resize(H,W), rotation, normalization &amp; convert to tensor</p>
<pre><code class="lang-auto">def check_accuracy(loader, model, device="cuda"):
    num_correct = 0
    num_pixels = 0
    dice_score = 0 #from [0,1] similar as F1 score
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            print('x {}, y {}'.format(x.shape,y.shape))
            y = y.to(device).unsqueeze(1) # y torch.Size([1, 180, 100, 3])
            preds = torch.sigmoid(model(x))
            preds = (preds &gt; 0.5).float() #preds torch.Size([1, 5, 180, 100])
            num_correct += (preds == y).sum()
            num_pixels += torch.numel(preds)
            dice_score += (2 * (preds * y).sum()) / (
                (preds + y).sum() + 1e-8
</code></pre>
<p>I already tried with <code>mask = np.array(Image.open(mask_path).convert("RGB"))</code> but it raised a conflict of tensor shapes  <code>preds torch.Size([1, 5, 180, 100]) vs y torch.Size([1, 180, 100, 3])</code> in <code>check_accuracy</code> function.</p>
<p>I used the training algo as below:</p>
<pre><code class="lang-auto">    model = UNET(in_channels=3, out_channels=5).to(DEVICE)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
</code></pre>
<p>This is the traceback error:</p>
<pre><code class="lang-auto">Traceback (most recent call last)
e:\04_Educations\MachineLearning\ML_Online\Pytorch\petroleum\train.py in &lt;cell line: 109&gt;()
    123         save_predictions_as_imgs(val_loader, model, folder="saved_images/", device=DEVICE)
    126 if __name__ == "__main__":
--&gt; 127     main()

e:\04_Educations\MachineLearning\ML_Online\Pytorch\petroleum\train.py in main()
     102 if LOAD_MODEL:
     103     load_checkpoint(torch.load("my_checkpoint.pth.tar"), model)
---&gt; 106 check_accuracy(val_loader, model, device=DEVICE)
     107 scaler = torch.cuda.amp.GradScaler()
     109 for epoch in range(NUM_EPOCHS):

e:\04_Educations\MachineLearning\ML_Online\Pytorch\utils.py in check_accuracy(loader, model, device)
     70 preds = (preds &gt; 0.5).float()
     71 print('check accuracy preds',preds.shape)
---&gt; 72 num_correct += (preds == y).sum()
     73 num_pixels += torch.numel(preds)
     74 dice_score += (2 * (preds * y).sum()) / (
     75     (preds + y).sum() + 1e-8
     76 )

RuntimeError: The size of tensor a (100) must match the size of tensor b (3) at non-singleton dimension 4
</code></pre>
<p>Please share any advice, thank you!</p>

ptrblck: <aside class="quote no-group" data-username="Bruce_Kieu" data-post="1" data-topic="146588">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bruce_kieu/48/44954_2.png" class="avatar"> Bruce_Kieu:</div>
<blockquote>
<p>but it raised a conflict of tensor shapes <code>preds torch.Size([1, 5, 180, 100]) vs y torch.Size([1, 180, 100, 3])</code> in <code>check_accuracy</code> function.</p>
</blockquote>
</aside>
<p>PyTorch uses the channels-first layout, so your target should have the shape <code>[1, 3, 180, 100]</code> if you want to use 3 channels.<br>
However, I’m not sure I fully understand the use case, since you won’t be able to compare the predictions for 5 classes with an RGB target in:</p>
<pre data-code-wrap="python"><code class="lang-python">num_correct += (preds == y).sum()
</code></pre>

Bruce_Kieu: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>,<br>
num_correct is used to compute accuracy of the training. so if the  <code>num_correct += (preds == y).sum()</code> not applicable, how should I revise it to use the target with 3 channels (RGB).<br>
I also tried to convert the input images to single channel <code>image = np.array(Image.open(img_path).convert("L"), dtype=np.float32)</code> but error persists there.<br>
Please let me share you a bit the context of tha: Actually my target is a 5 geology layers model:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/b/ebaefe1182a4766e787f5fbc7ab50be2656a34de.png" data-download-href="https://discuss.pytorch.org/uploads/default/ebaefe1182a4766e787f5fbc7ab50be2656a34de" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/b/ebaefe1182a4766e787f5fbc7ab50be2656a34de_2_282x500.png" alt="image" data-base62-sha1="xCX7L7gOBArm0zxkXbz67lltB38" width="282" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/b/ebaefe1182a4766e787f5fbc7ab50be2656a34de_2_282x500.png, https://discuss.pytorch.org/uploads/default/original/3X/e/b/ebaefe1182a4766e787f5fbc7ab50be2656a34de.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/b/ebaefe1182a4766e787f5fbc7ab50be2656a34de.png 2x" data-dominant-color="3C5871"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">308×545 4.53 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div> That’s why i set out_channel=5. My goal is using Unet model derives feature from the inputs as here<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/9/9/99e5e276f658d1f7e808fa714dd0904808b62c5b.png" alt="image" data-base62-sha1="lXrzRQ0PlbYBvgBIWybw9qb06eL" width="332" height="196"> to predict output as similar to the target as possible. I have about 100 image &amp; mask pairs, I can increase the number of data if needed.<br>
I’m new on this field, hope to get more advice, thank you!</p>

ptrblck: <aside class="quote no-group" data-username="Bruce_Kieu" data-post="3" data-topic="146588">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bruce_kieu/48/44954_2.png" class="avatar"> Bruce_Kieu:</div>
<blockquote>
<p>num_correct is used to compute accuracy of the training. so if the <code>num_correct += (preds == y).sum()</code> not applicable, how should I revise it to use the target with 3 channels (RGB).</p>
</blockquote>
</aside>
<p>Based on your shared code it seems you are using <code>nn.CrossEntropyLoss</code> which expects the targets to contain class indices in the default case.<br>
If you’ve already created these targets, you could use them to calculate the accuracy by comparing them against <code>preds</code>.<br>
However, since <code>nn.CrossEntropyLoss</code> expects the model to output raw logits, this line of code also looks wrong:</p>
<pre data-code-wrap="python"><code class="lang-python">preds = (preds &gt; 0.5).float()
</code></pre>
<p>since you are comparing against a probability threshold while <code>preds = torch.argmax(output, dim=1)</code> is usually used for a multi-class segmentation.</p>

Bruce_Kieu: <p>Thanks ptrblck, let me dive deeper to that</p>