Isaac_Kargar:

For the first question, if you put a sample code, we can help you more. It is because of shape mismatching


Certainly! I will put here the code of the data loader and the train step.
In the Data Loader, L channel and a,b channels must be returned for the training step.
#####################################################
################## DATA LOADER #####################
#####################################################

from torchvision import datasets, transforms
from torch.utils.data import Dataset
from skimage.color import rgb2lab, rgb2gray
from skimage import io
import torch.utils.data as data
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
import os

scale_transform = transforms.Compose([
    transforms.Resize((256,256),2),
    #transforms.RandomCrop(224),
    #transforms.ToTensor()
])

class CustomDataset(Dataset):
    """Custom Dataset."""

    def __init__(self, root_dir, transform=None):
        """
        Args:
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.root_dir = root_dir
        self.transform = transform
        self.file_list=os.listdir(root_dir)
        self.tensor_to_PIL = transforms.ToPILImage()

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

    def __getitem__(self, idx):
        img = Image.open(self.root_dir+'/'+self.file_list[idx]) # Read the image
        
        if self.transform is not None:
            img_rgb_resized = transforms.Resize((64,64),2)(img)
            img_rgb_ori = np.array(img)
            
            img_lab_ori = rgb2lab(img_rgb_ori) # Convert to CIE Lab color space
            img_lab_resized = rgb2lab(img_rgb_resized) # Convert to CIE Lab color space
            
            img_rgb_transposed = img_rgb_ori.transpose(2, 0, 1) #(C, W, H)    
            img_lab_transposed = img_lab_ori.transpose(2, 0, 1) #(C, W, H)
            img_lab_resized_transposed = img_lab_resized.transpose(2, 0, 1) #(C, W, H)
            
            img_l = (np.round(img_lab_transposed[0,:,:])).astype(np.int) # L channel
            img_l = self.tensor_to_PIL(img_l) # Convert to PIL object image to be possible apply the following transform
            img_l_resized = self.transform(img_l)
            img_l_resized = np.array(img_l_resized) # Convert to numpy array
            img_l_resized = torch.from_numpy(img_l_resized) # Convert to torch tensor
            
            img_ab_resized = (np.round(img_lab_resized_transposed[1:3, :, :])).astype(np.int) # (a,b) channels with int intensity values            
            img_ab_resized = np.array(img_ab_resized) # Convert to numpy array            
            img_ab_resized = torch.from_numpy(img_ab_resized) # Convert to torch tensor

            filename = self.root_dir+'/'+self.file_list[idx]
            
            return img_l_resized, img_ab_resized, filename # img_l_resized -> 1x256x256 and img_ab_resized -> 2x64x64

###################################################
################## TRAIN STEP #####################
###################################################
import os
import torch
import argparse
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
import numpy as np
from training_layers import PriorBoostLayer, NNEncLayer, ClassRebalanceMultLayer, NonGrayMaskLayer
from data_loader import TrainImageFolder, CustomDataset
from model import Color_model
np.set_printoptions(threshold=np.inf)

original_transform = transforms.Compose([
    transforms.Resize((256,256),2),    
])


def main(args):
    
    train_set = CustomDataset(args.image_dir, original_transform)

    data_loader = torch.utils.data.DataLoader(train_set, batch_size = args.batch_size, shuffle = True, num_workers = args.num_workers)

    model = Color_model().cuda()
    criterion = nn.CrossEntropyLoss().cuda()
    params = list(model.parameters())
    optimizer = torch.optim.Adam(params, lr = args.learning_rate)

    encode_ab_layer = NNEncLayer()   

    #######################
    ### Train the model ###
    #######################

    total_step = len(data_loader)
    
    for epoch in range(args.num_epochs):
        for i, (images, img_ab, filename) in enumerate(data_loader):
            #print(filename)
            images = images.unsqueeze(1).float().cuda()
            img_ab = img_ab.float()
                        
            encode_ab, max_encode_ab = encode_ab_layer.forward(img_ab)
            encode_ab = torch.from_numpy(encode_ab).long().cuda()

            targets=torch.Tensor(max_encode_ab).long().cuda()
            print(images.shape)
            outputs = model(images)
            
            loss=criterion(outputs,targets)
            
            model.zero_grad()
            
            loss.backward()
            optimizer.step()

            # Print log info
            if i % args.log_step == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch, args.num_epochs, i, total_step, loss.item()))

            # Save the model checkpoints

            if epoch == 4:
                torch.save(model.state_dict(), os.path.join(args.model_path, 'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))

            
            
            

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_dir', type = str, default = 'C:\\Users\\Matheus Santos\\Desktop\\Colorization-Epiphqny\\Dataset\\train\\images', help = 'directory for resized images')
    parser.add_argument('--model_path', type = str, default = 'C:\\Users\\Matheus Santos\\Desktop\\Colorization-Epiphqny\\models', help = 'path for saving trained models')
    parser.add_argument('--crop_size', type = int, default = 224, help = 'size for randomly cropping images')
    parser.add_argument('--log_step', type = int, default = 1, help = 'step size for prining log info')
    parser.add_argument('--save_step', type = int, default = 5, help = 'step size for saving trained models')

    # Model parameters
    parser.add_argument('--num_epochs', type = int, default = 5)
    parser.add_argument('--batch_size', type = int, default = 10)
    parser.add_argument('--num_workers', type = int, default = 8)
    parser.add_argument('--learning_rate', type = float, default = 1e-3)
    args = parser.parse_args()
    print(args)
    main(args)


Any doubts about the code, ask me please