Hi K. Frank, thanks so much for your notes.
Sorry, I forgot to enter the training code, it’s below:
import torch
import albumentations as A
from albumentations.pytorch import ToTensorV2
from tqdm import tqdm
import torch.nn as nn
import torch.optim as optim
from model import UNet
from utils import (
    load_checkpoint,
    save_checkpoint,
    get_loaders,
    check_accuracy,
    save_predictions_as_imgs,
)
print(f'PyTorch: {torch.__version__}')
print(f'CUDA: {torch.version.cuda}')

if torch.cuda.is_available() == True:
   DEVICE = "cuda"
   print('GPU')
   #torch.manual_seed(12)
   #torch.cuda.manual_seed(12)
   #np.random.seed(12)
else: 
   DEVICE = "cpu"
   print('CPU')

if(torch.backends.cudnn.is_available() == True):
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    print('cuDNN activated')
else:
    print('cuDNN not available!')

# Hyperparameters etc.
LEARNING_RATE = 3e-4
BATCH_SIZE = 1
NUM_EPOCHS = 10
NUM_WORKERS = 0
IMAGE_HEIGHT = 160
IMAGE_WIDTH = 240
PIN_MEMORY = True
LOAD_MODEL = False
TRAIN_IMG_DIR = "DATASET/train/imagens/"
TRAIN_MASK_DIR = "DATASET/train/mascaras/"
VAL_IMG_DIR = "DATASET/val/imagens/"
VAL_MASK_DIR = "DATASET/val/mascaras/"

#torch.cuda.empty_cache()

def train_fn(loader, model, optimizer, loss_fn, scaler):
    loop = tqdm(loader)

    for batch_idx, (data, targets) in enumerate(loop):
        data = data.to(device=DEVICE)
        targets = targets.long().squeeze(1).to(device=DEVICE)

        # forward
        with torch.cuda.amp.autocast():
            predictions = model(data)
            loss = loss_fn(predictions, targets)
            
        # backward
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        # update tqdm loop
        loop.set_postfix(loss=loss.item())


def main():
    train_transform = A.Compose(
        [
            A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
            A.Rotate(limit=35, p=1.0),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.1),
            A.Normalize(
                mean=[0.0, 0.0, 0.0],
                std=[1.0, 1.0, 1.0],
                max_pixel_value=255.0,
            ),
            ToTensorV2(),
        ],
    )

    val_transforms = A.Compose(
        [
            A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH),
            A.Normalize(
                mean=[0.0, 0.0, 0.0],
                std=[1.0, 1.0, 1.0],
                max_pixel_value=255.0,
            ),
            ToTensorV2(),
        ],
    )

    
    model = UNet(21).to(DEVICE)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    train_loader, val_loader = get_loaders(
        TRAIN_IMG_DIR,
        TRAIN_MASK_DIR,
        VAL_IMG_DIR,
        VAL_MASK_DIR,
        BATCH_SIZE,
        train_transform,
        val_transforms,
        NUM_WORKERS,
        PIN_MEMORY,
    )

    if LOAD_MODEL:
        load_checkpoint(torch.load("my_checkpoint.pth.tar"), model)


    check_accuracy(val_loader, model, device=DEVICE)
    scaler = torch.cuda.amp.GradScaler()

    for epoch in range(NUM_EPOCHS):
        train_fn(train_loader, model, optimizer, loss_fn, scaler)

        # save model
        checkpoint = {
            "state_dict": model.state_dict(),
            "optimizer":optimizer.state_dict(),
        }
        save_checkpoint(checkpoint)

        # check accuracy
        check_accuracy(val_loader, model, device=DEVICE)

        # print some examples to a folder
        save_predictions_as_imgs(
            val_loader, model, folder="Results\\", device=DEVICE
        )


if __name__ == "__main__":
    main()


My code may seem meaningless at times, sorry about that, but I’m a beginner. I will try to explain some points better.


I was also surprised that the training code didn’t report a bug. But if I apply argmax function to predictions and target to remove dimension 1, the error in check_accuracy function disappears. But then it causes a similar error in the training code.


I forgot to comment, but I convert the images to RGB, because if I don’t, the colors of the masks are changed. When this happens, some different classes are given the same color. I don’t know why this happens. I’ll leave the dataset code here:


import os
from PIL import Image
from torch.utils.data import Dataset
import numpy as np
import time
import matplotlib
import matplotlib.pyplot as plt
import cv2
import torch

class MeuDataset(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].replace(".jpg", ".png"))
        image = np.array(Image.open(img_path).convert("RGB"))
        mask = np.array(Image.open(mask_path).convert("RGB"), dtype=np.uint8)
        #mask = np.transpose(mask,(2, 0, 1))        
        #debug
        #print(f'Normal: {mask.shape}')
        #plt.imshow(mask)
        #plt.show()
        #mask = mask/classes
        #mask[mask == 255.0] = 1.0

        if self.transform is not None:
            augmentations = self.transform(image=image, mask=mask)
            image = augmentations["image"]
            mask = augmentations["mask"]

        return image, mask


Target formatting is really my biggest issue here, as you said. I’ve seen in other threads that we must convert RGB colors to class indices, but I have no idea how to do that. I’ve even tried it, but nothing has worked yet.

After changing the accuracy calculation as you suggested, I got the following error:
1only batches of spatial targets supported (3D tensors) but got targets of size: : [1, 160, 240, 3].
In this case, swapping the dimensions as required by PyTorch and removing the dimension corresponding to 3 would be sufficient?