hey
i tried creating a unet-u2net architecture

here's the complete code
help me with the errors

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, random_split
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm  # For progress bars
import warnings
import copy
​
warnings.filterwarnings("ignore")
add Codeadd Markdown
def encode_segmap(mask):
    """
    Converts an RGB segmentation mask to a 2D tensor of class indices.
    
    Args:
        mask (PIL Image): RGB segmentation mask.
        
    Returns:
        torch.Tensor: 2D tensor with class indices.
    """
    mask = np.array(mask)
    # Initialize the encoded mask with zeros (Background)
    encoded_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int64)
    
    # Define the color codes for each class
    # Adjust these values based on the dataset's encoding
    background = [0, 0, 0]
    pet = [128, 0, 0]
    border = [255, 0, 0]
    
    # Assign class indices based on color
    encoded_mask[np.all(mask == background, axis=-1)] = 0
    encoded_mask[np.all(mask == pet, axis=-1)] = 1
    encoded_mask[np.all(mask == border, axis=-1)] = 2
    
    return torch.from_numpy(encoded_mask)
​
# Image transformations
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
])
​
# Mask transformations
mask_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.Lambda(lambda mask: encode_segmap(mask)),
])
add Codeadd Markdown
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

mask_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

# 3. Load the Oxford-IIIT Pet Dataset with 'trainval' Split
trainval_dataset = datasets.OxfordIIITPet(
    root='./data',
    split='trainval',
    target_types='segmentation',
    transform=transform,
    target_transform=mask_transform,
    download=True
)

test_dataset = datasets.OxfordIIITPet(
    root='./data',
    split='test',
    target_types='segmentation',
    transform=transform,
    target_transform=mask_transform,
    download=True
)

# 4. Create Training and Validation Splits from 'trainval'
validation_size = 0.2  # 20% for validation
train_size = int((1 - validation_size) * len(trainval_dataset))
val_size = len(trainval_dataset) - train_size

train_dataset, val_dataset = random_split(trainval_dataset, [train_size, val_size])

print(f"Total 'trainval' samples: {len(trainval_dataset)}")
print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")
print(f"Test samples: {len(test_dataset)}\n")

# 5. Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=2)

# 6. Visualize a Sample from the Training Set
def visualize_sample(dataset, idx):
    image, mask = dataset[idx]
    image = image.permute(1, 2, 0).numpy()
    mask = mask.squeeze().numpy()
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title('Input Image')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.imshow(mask, cmap='gray')
    plt.title('Segmentation Mask')
    plt.axis('off')
    plt.show()

# Example visualization
visualize_sample(train_dataset, 0)
# Load the Oxford-IIIT Pet Dataset with 'trainval' Split
trainval_dataset = datasets.OxfordIIITPet(
    root='./data',
    split='trainval',  # Correct split
    target_types='segmentation',
    transform=transform,
    target_transform=mask_transform,
    download=True
)
​
test_dataset = datasets.OxfordIIITPet(
    root='./data',
    split='test',
    target_types='segmentation',
    transform=transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
    ]),  # No augmentation for test
    target_transform=transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.Lambda(lambda mask: encode_segmap(mask)),
    ]),
    download=True
)
​
# Create Training and Validation Splits from 'trainval'
validation_size = 0.2  # 20% for validation
train_size = int((1 - validation_size) * len(trainval_dataset))
val_size = len(trainval_dataset) - train_size
​
torch.manual_seed(42)  # Set a seed for reproducibility
train_dataset, val_dataset = random_split(trainval_dataset, [train_size, val_size])
​
print(f"Total 'trainval' samples: {len(trainval_dataset)}")
print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")
print(f"Test samples: {len(test_dataset)}\n")
Total 'trainval' samples: 3680
Training samples: 2944
Validation samples: 736
Test samples: 3669

add Codeadd Markdown
class DoubleConv(nn.Module):
    """[Conv2d => BatchNorm => ReLU] x 2"""

    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),

            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.double_conv(x)

class AttentionBlock(nn.Module):
    """Spatial and Channel Attention"""

    def __init__(self, in_channels, gating_channels, inter_channels=None):
        super(AttentionBlock, self).__init__()
        
        if inter_channels is None:
            inter_channels = in_channels // 2
        
        self.W_g = nn.Sequential(
            nn.Conv2d(gating_channels, inter_channels, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(inter_channels)
        )

        self.W_x = nn.Sequential(
            nn.Conv2d(in_channels, inter_channels, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(inter_channels)
        )

        self.psi = nn.Sequential(
            nn.Conv2d(inter_channels, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x, g):
        """
            x: input feature (from encoder)
            g: gating signal (from bottle neck)
        """
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi

class NestedUNetBlock(nn.Module):
    """Nested U-Net Block inspired by U²Net"""

    def __init__(self, in_channels, out_channels, sub_nested_depth=2):
        super(NestedUNetBlock, self).__init__()
        self.sub_nested_depth = sub_nested_depth
        self.main_conv = DoubleConv(in_channels, out_channels)
        
        self.nested_convs = nn.ModuleList()
        self.attentions = nn.ModuleList()
        
        for _ in range(sub_nested_depth):
            self.nested_convs.append(DoubleConv(out_channels, out_channels))
            self.attentions.append(AttentionBlock(out_channels, out_channels))

    def forward(self, x, gating_signal=None):
        """
            x: input feature
            gating_signal: feature from bottleneck for attention (if any)
        """
        out = self.main_conv(x)
        for i in range(self.sub_nested_depth):
            if gating_signal is not None and i == self.sub_nested_depth - 1:
                attn = self.attentions[i](out, gating_signal)
                out = self.nested_convs[i](attn)
            else:
                out = self.nested_convs[i](out)
        return out

class MultiScaleFusion(nn.Module):
    """Feature Pyramid Network-like Fusion"""

    def __init__(self, in_channels, out_channels):
        super(MultiScaleFusion, self).__init__()
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        """
            x1: higher resolution feature
            x2: lower resolution feature
        """
        x2 = F.interpolate(x2, size=x1.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat([x1, x2], dim=1)
        x = self.conv(x)
        return x
# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False, num_workers=2)
add Codeadd Markdown
class U2UNet(nn.Module):
    def __init__(self, n_channels=3, n_classes=3, base_filters=64):
        super(U2UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        filters = base_filters

        # Encoder
        self.encoder1 = NestedUNetBlock(n_channels, filters)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder2 = NestedUNetBlock(filters, filters*2)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder3 = NestedUNetBlock(filters*2, filters*4)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder4 = NestedUNetBlock(filters*4, filters*8)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Bottleneck
        self.bottleneck = NestedUNetBlock(filters*8, filters*16, sub_nested_depth=2)

        # Decoder
        self.up4 = nn.ConvTranspose2d(filters*16, filters*8, kernel_size=2, stride=2)
        self.attn4 = AttentionBlock(filters*8, filters*8)
        self.decoder4 = NestedUNetBlock(filters*16, filters*8)

        self.up3 = nn.ConvTranspose2d(filters*8, filters*4, kernel_size=2, stride=2)
        self.attn3 = AttentionBlock(filters*4, filters*4)
        self.decoder3 = NestedUNetBlock(filters*8, filters*4)

        self.up2 = nn.ConvTranspose2d(filters*4, filters*2, kernel_size=2, stride=2)
        self.attn2 = AttentionBlock(filters*2, filters*2)
        self.decoder2 = NestedUNetBlock(filters*4, filters*2)

        self.up1 = nn.ConvTranspose2d(filters*2, filters, kernel_size=2, stride=2)
        self.attn1 = AttentionBlock(filters, filters)
        self.decoder1 = NestedUNetBlock(filters*2, filters)

        # Final Convolution
        self.final_conv = nn.Conv2d(filters, n_classes, kernel_size=1)
    
    def forward(self, x):
        # Encoder path
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))
        bottleneck = self.bottleneck(self.pool4(enc4))

        # Decoder path
        dec4 = self.up4(bottleneck)
        enc4_attn = self.attn4(enc4, dec4)
        dec4 = torch.cat((enc4_attn, dec4), dim=1)
        dec4 = self.decoder4(dec4)

        dec3 = self.up3(dec4)
        enc3_attn = self.attn3(enc3, dec3)
        dec3 = torch.cat((enc3_attn, dec3), dim=1)
        dec3 = self.decoder3(dec3)

        dec2 = self.up2(dec3)
        enc2_attn = self.attn2(enc2, dec2)
        dec2 = torch.cat((enc2_attn, dec2), dim=1)
        dec2 = self.decoder2(dec2)

        dec1 = self.up1(dec2)
        enc1_attn = self.attn1(enc1, dec1)
        dec1 = torch.cat((enc1_attn, dec1), dim=1)
        dec1 = self.decoder1(dec1)

        # Final output without activation
        out = self.final_conv(dec1)  # Raw logits
        return out
def visualize_sample(dataset, idx):
    """
    Visualizes an image and its corresponding segmentation mask.
    
    Args:
        dataset (Dataset): The dataset to visualize from.
        idx (int): Index of the sample to visualize.
    """
    image, mask = dataset[idx]
    image = image.permute(1, 2, 0).numpy()
    mask = mask.numpy()
    
    # Define a color map for visualization
    color_map = {
        0: [0, 0, 0],        # Background - Black
        1: [255, 255, 255],  # Pet - White
        2: [128, 128, 128],  # Border - Gray
    }
    
    mask_color = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    for class_idx, color in color_map.items():
        mask_color[mask == class_idx] = color
    
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 3, 1)
    plt.imshow(image)
    plt.title('Input Image')
    plt.axis('off')
    
    plt.subplot(1, 3, 2)
    plt.imshow(mask, cmap='gray')
    plt.title('Segmentation Mask (Class Indices)')
    plt.axis('off')
    
    plt.subplot(1, 3, 3)
    plt.imshow(mask_color)
    plt.title('Segmentation Mask (Colored)')
    plt.axis('off')
    
    plt.show()
​
# Example visualization
visualize_sample(train_dataset, 0)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[16], line 43
     40     plt.show()
     42 # Example visualization
---> 43 visualize_sample(train_dataset, 0)

Cell In[16], line 9, in visualize_sample(dataset, idx)
      1 def visualize_sample(dataset, idx):
      2     """
      3     Visualizes an image and its corresponding segmentation mask.
      4     
   (...)
      7         idx (int): Index of the sample to visualize.
      8     """
----> 9     image, mask = dataset[idx]
     10     image = image.permute(1, 2, 0).numpy()
     11     mask = mask.numpy()

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataset.py:412, in Subset.__getitem__(self, idx)
    410 if isinstance(idx, list):
    411     return self.dataset[[self.indices[i] for i in idx]]
--> 412 return self.dataset[self.indices[idx]]

File /opt/conda/lib/python3.10/site-packages/torchvision/datasets/oxford_iiit_pet.py:116, in OxfordIIITPet.__getitem__(self, idx)
    113     target = tuple(target)
    115 if self.transforms:
--> 116     image, target = self.transforms(image, target)
    118 return image, target

File /opt/conda/lib/python3.10/site-packages/torchvision/datasets/vision.py:97, in StandardTransform.__call__(self, input, target)
     95     input = self.transform(input)
     96 if self.target_transform is not None:
---> 97     target = self.target_transform(target)
     98 return input, target

File /opt/conda/lib/python3.10/site-packages/torchvision/transforms/transforms.py:95, in Compose.__call__(self, img)
     93 def __call__(self, img):
     94     for t in self.transforms:
---> 95         img = t(img)
     96     return img

File /opt/conda/lib/python3.10/site-packages/torchvision/transforms/transforms.py:479, in Lambda.__call__(self, img)
    478 def __call__(self, img):
--> 479     return self.lambd(img)

Cell In[13], line 39, in <lambda>(mask)
     29 transform = transforms.Compose([
     30     transforms.Resize((256, 256)),
     31     transforms.RandomHorizontalFlip(),
     32     transforms.RandomRotation(10),
     33     transforms.ToTensor(),
     34 ])
     36 # Mask transformations
     37 mask_transform = transforms.Compose([
     38     transforms.Resize((256, 256)),
---> 39     transforms.Lambda(lambda mask: encode_segmap(mask)),
     40 ])

Cell In[13], line 22, in encode_segmap(mask)
     19 border = [255, 0, 0]
     21 # Assign class indices based on color
---> 22 encoded_mask[np.all(mask == background, axis=-1)] = 0
     23 encoded_mask[np.all(mask == pet, axis=-1)] = 1
     24 encoded_mask[np.all(mask == border, axis=-1)] = 2

ValueError: operands could not be broadcast together with shapes (256,256) (3,) 
add Codeadd Markdown
from tqdm import tqdm  # For progress bars

def train_model(model, dataloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    loop = tqdm(dataloader, desc='Training', leave=False)
    for images, masks in loop:
        images = images.to(device)
        masks = masks.to(device).long()  # Ensure masks are of type long for CrossEntropyLoss

        optimizer.zero_grad()
        outputs = model(images)  # Raw logits
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * images.size(0)

        preds = torch.argmax(outputs, dim=1)
        correct += (preds == masks).sum().item()
        total += masks.numel()

        loop.set_postfix(loss=loss.item())

    epoch_loss = running_loss / len(dataloader.dataset)
    epoch_acc = 100 * correct / total

    return epoch_loss, epoch_acc

def evaluate_model(model, dataloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    loop = tqdm(dataloader, desc='Evaluating', leave=False)
    with torch.no_grad():
        for images, masks in loop:
            images = images.to(device)
            masks = masks.to(device).long()

            outputs = model(images)
            loss = criterion(outputs, masks)

            running_loss += loss.item() * images.size(0)

            preds = torch.argmax(outputs, dim=1)
            correct += (preds == masks).sum().item()
            total += masks.numel()

            loop.set_postfix(loss=loss.item())

    epoch_loss = running_loss / len(dataloader.dataset)
    epoch_acc = 100 * correct / total

    return epoch_loss, epoch_acc
class DoubleConv(nn.Module):
    """[Conv2d => BatchNorm => ReLU] x 2"""
​
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
​
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )
​
    def forward(self, x):
        return self.double_conv(x)
​
class AttentionBlock(nn.Module):
    """Spatial and Channel Attention"""
​
    def __init__(self, in_channels, gating_channels, inter_channels=None):
        super(AttentionBlock, self).__init__()
        
        if inter_channels is None:
            inter_channels = in_channels // 2
        
        self.W_g = nn.Sequential(
            nn.Conv2d(gating_channels, inter_channels, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(inter_channels)
        )
​
        self.W_x = nn.Sequential(
            nn.Conv2d(in_channels, inter_channels, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(inter_channels)
        )
​
        self.psi = nn.Sequential(
            nn.Conv2d(inter_channels, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )
​
        self.relu = nn.ReLU(inplace=True)
​
    def forward(self, x, g):
        """
            x: input feature (from encoder)
            g: gating signal (from bottleneck)
        """
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi
​
class NestedUNetBlock(nn.Module):
    """Nested U-Net Block inspired by U²Net"""
​
    def __init__(self, in_channels, out_channels, sub_nested_depth=2):
        super(NestedUNetBlock, self).__init__()
        self.sub_nested_depth = sub_nested_depth
        self.main_conv = DoubleConv(in_channels, out_channels)
        
        self.nested_convs = nn.ModuleList()
        self.attentions = nn.ModuleList()
        
        for _ in range(sub_nested_depth):
            self.nested_convs.append(DoubleConv(out_channels, out_channels))
            self.attentions.append(AttentionBlock(out_channels, out_channels))
​
    def forward(self, x, gating_signal=None):
        """
            x: input feature
            gating_signal: feature from bottleneck for attention (if any)
        """
        out = self.main_conv(x)
        for i in range(self.sub_nested_depth):
            if gating_signal is not None and i == self.sub_nested_depth - 1:
                attn = self.attentions[i](out, gating_signal)
                out = self.nested_convs[i](attn)
            else:
                out = self.nested_convs[i](out)
        return out
​
class MultiScaleFusion(nn.Module):
    """Feature Pyramid Network-like Fusion"""
​
    def __init__(self, in_channels, out_channels):
        super(MultiScaleFusion, self).__init__()
        self.conv = DoubleConv(in_channels, out_channels)
​
    def forward(self, x1, x2):
        """
            x1: higher resolution feature
            x2: lower resolution feature
        """
        x2 = F.interpolate(x2, size=x1.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat([x1, x2], dim=1)
        x = self.conv(x)
        return x
​
class U2UNet(nn.Module):
    def __init__(self, n_channels=3, n_classes=3, base_filters=64):
        super(U2UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        filters = base_filters
​
        # Encoder
        self.encoder1 = NestedUNetBlock(n_channels, filters)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
​
        self.encoder2 = NestedUNetBlock(filters, filters*2)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
​
        self.encoder3 = NestedUNetBlock(filters*2, filters*4)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
​
        self.encoder4 = NestedUNetBlock(filters*4, filters*8)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
​
        # Bottleneck
        self.bottleneck = NestedUNetBlock(filters*8, filters*16, sub_nested_depth=2)
​
        # Decoder
        self.up4 = nn.ConvTranspose2d(filters*16, filters*8, kernel_size=2, stride=2)
        self.attn4 = AttentionBlock(filters*8, filters*8)
        self.decoder4 = NestedUNetBlock(filters*16, filters*8)
​
        self.up3 = nn.ConvTranspose2d(filters*8, filters*4, kernel_size=2, stride=2)
        self.attn3 = AttentionBlock(filters*4, filters*4)
        self.decoder3 = NestedUNetBlock(filters*8, filters*4)
​
        self.up2 = nn.ConvTranspose2d(filters*4, filters*2, kernel_size=2, stride=2)
        self.attn2 = AttentionBlock(filters*2, filters*2)
        self.decoder2 = NestedUNetBlock(filters*4, filters*2)
​
        self.up1 = nn.ConvTranspose2d(filters*2, filters, kernel_size=2, stride=2)
        self.attn1 = AttentionBlock(filters, filters)
        self.decoder1 = NestedUNetBlock(filters*2, filters)
​
        # Final Convolution
        self.final_conv = nn.Conv2d(filters, n_classes, kernel_size=1)
    
    def forward(self, x):
        # Encoder path
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))
        bottleneck = self.bottleneck(self.pool4(enc4))
​
        # Decoder path
        dec4 = self.up4(bottleneck)
        enc4_attn = self.attn4(enc4, dec4)
        dec4 = torch.cat((enc4_attn, dec4), dim=1)
        dec4 = self.decoder4(dec4)
​
        dec3 = self.up3(dec4)
        enc3_attn = self.attn3(enc3, dec3)
        dec3 = torch.cat((enc3_attn, dec3), dim=1)
        dec3 = self.decoder3(dec3)
​
        dec2 = self.up2(dec3)
        enc2_attn = self.attn2(enc2, dec2)
        dec2 = torch.cat((enc2_attn, dec2), dim=1)
        dec2 = self.decoder2(dec2)
​
        dec1 = self.up1(dec2)
        enc1_attn = self.attn1(enc1, dec1)
        dec1 = torch.cat((enc1_attn, dec1), dim=1)
        dec1 = self.decoder1(dec1)
​
        # Final output without activation
        out = self.final_conv(dec1)  # Raw logits
        return out
add Codeadd Markdown
# Initialize model, loss function, optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = U2UNet(n_channels=3, n_classes=3, base_filters=64).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# Training parameters
num_epochs = 25

# To store losses and accuracies
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []

for epoch in range(num_epochs):
    train_loss, train_acc = train_model(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate_model(model, val_loader, criterion, device)

    train_losses.append(train_loss)
    train_accuracies.append(train_acc)
    val_losses.append(val_loss)
    val_accuracies.append(val_acc)

    print(f"Epoch [{epoch+1}/{num_epochs}]")
    print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%\n")
def train_model(model, dataloader, criterion, optimizer, device):
    """
    Trains the model for one epoch.
    """
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
​
    loop = tqdm(dataloader, desc='Training', leave=False)
    for images, masks in loop:
        images = images.to(device)
        masks = masks.to(device).long()  # Ensure masks are of type long for CrossEntropyLoss
​
        optimizer.zero_grad()
        outputs = model(images)  # Raw logits
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
​
        running_loss += loss.item() * images.size(0)
​
        preds = torch.argmax(outputs, dim=1)
        correct += (preds == masks).sum().item()
        total += masks.numel()
​
        loop.set_postfix(loss=loss.item())
​
    epoch_loss = running_loss / len(dataloader.dataset)
    epoch_acc = 100 * correct / total
​
    return epoch_loss, epoch_acc
​
def evaluate_model(model, dataloader, criterion, device):
    """
    Evaluates the model on the validation set.
    """
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
​
    loop = tqdm(dataloader, desc='Evaluating', leave=False)
    with torch.no_grad():
        for images, masks in loop:
            images = images.to(device)
            masks = masks.to(device).long()
​
            outputs = model(images)
            loss = criterion(outputs, masks)
​
            running_loss += loss.item() * images.size(0)
​
            preds = torch.argmax(outputs, dim=1)
            correct += (preds == masks).sum().item()
            total += masks.numel()
​
            loop.set_postfix(loss=loss.item())
​
    epoch_loss = running_loss / len(dataloader.dataset)
    epoch_acc = 100 * correct / total
​
    return epoch_loss, epoch_acc
add Codeadd Markdown
# Plotting Losses
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(range(1, num_epochs+1), train_losses, label='Train Loss')
plt.plot(range(1, num_epochs+1), val_losses, label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss Over Epochs')
plt.legend()

# Plotting Accuracies
plt.subplot(1, 2, 2)
plt.plot(range(1, num_epochs+1), train_accuracies, label='Train Accuracy')
plt.plot(range(1, num_epochs+1), val_accuracies, label='Val Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.title('Accuracy Over Epochs')
plt.legend()

plt.show()
class EarlyStopping:
    def __init__(self, patience=5, verbose=False, delta=0):
        """
        Args:
            patience (int): How long to wait after last time validation loss improved.
            verbose (bool): If True, prints a message for each validation loss improvement.
            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
        """
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_loss = None
        self.early_stop = False
        self.delta = delta
        self.best_model_wts = None
​
    def __call__(self, val_loss, model):
        if self.best_loss is None:
            self.best_loss = val_loss
            self.best_model_wts = copy.deepcopy(model.state_dict())
            if self.verbose:
                print(f"Initial validation loss: {self.best_loss:.4f}")
        elif val_loss > self.best_loss - self.delta:
            self.counter += 1
            if self.verbose:
                print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.best_model_wts = copy.deepcopy(model.state_dict())
            self.counter = 0
            if self.verbose:
                print(f"Validation loss decreased to {self.best_loss:.4f}. Resetting counter.")
add Codeadd Markdown
arrow_upwardarrow_downwarddelete
def visualize_predictions(model, dataloader, device, num_samples=5):
    model.eval()
    samples = 0
    plt.figure(figsize=(15, num_samples * 5))
    
    with torch.no_grad():
        for images, masks in dataloader:
            images = images.to(device)
            masks = masks.to(device)

            outputs = model(images)
            preds = (outputs > 0.5).float()

            for i in range(images.size(0)):
                if samples >= num_samples:
                    return
                image = images[i].cpu().permute(1, 2, 0).numpy()
                mask = masks[i].squeeze().cpu().numpy()
                pred = preds[i].squeeze().cpu().numpy()

                plt.subplot(num_samples, 3, samples * 3 + 1)
                plt.imshow(image)
                plt.title('Input Image')
                plt.axis('off')

                plt.subplot(num_samples, 3, samples * 3 + 2)
                plt.imshow(mask, cmap='gray')
                plt.title('True Mask')
                plt.axis('off')

                plt.subplot(num_samples, 3, samples * 3 + 3)
                plt.imshow(pred, cmap='gray')
                plt.title('Predicted Mask')
                plt.axis('off')

                samples += 1

visualize_predictions(model, test_loader, device, num_samples=5)
# Initialize model, loss function, optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Training on device: {device}")
​
model = U2UNet(n_channels=3, n_classes=3, base_filters=64).to(device)
print("Model initialized.")
​
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
print("Loss function and optimizer set.")
​
# Training parameters
num_epochs = 25
​
# Early Stopping
early_stopping = EarlyStopping(patience=7, verbose=True)
​
# To store losses and accuracies
train_losses = []
train_accuracies = []
val_losses = []
val_accuracies = []
​
for epoch in range(num_epochs):
    print(f"Starting Epoch {epoch+1}/{num_epochs}")
    train_loss, train_acc = train_model(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = evaluate_model(model, val_loader, criterion, device)
​
    train_losses.append(train_loss)
    train_accuracies.append(train_acc)
    val_losses.append(val_loss)
    val_accuracies.append(val_acc)
​
    print(f"Epoch [{epoch+1}/{num_epochs}] Summary:")
    print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%\n")
    
    # Call EarlyStopping
    early_stopping(val_loss, model)
    
    if early_stopping.early_stop:
        print("Early stopping triggered. Training halted.")
        break
​
# Load the best model weights
if early_stopping.best_model_wts:
    model.load_state_dict(early_stopping.best_model_wts)
    print("Loaded best model weights.")
Training on device: cuda
Model initialized.
Loss function and optimizer set.
Starting Epoch 1/25
                                                 
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[20], line 26
     24 for epoch in range(num_epochs):
     25     print(f"Starting Epoch {epoch+1}/{num_epochs}")
---> 26     train_loss, train_acc = train_model(model, train_loader, criterion, optimizer, device)
     27     val_loss, val_acc = evaluate_model(model, val_loader, criterion, device)
     29     train_losses.append(train_loss)

Cell In[18], line 11, in train_model(model, dataloader, criterion, optimizer, device)
      8 total = 0
     10 loop = tqdm(dataloader, desc='Training', leave=False)
---> 11 for images, masks in loop:
     12     images = images.to(device)
     13     masks = masks.to(device).long()  # Ensure masks are of type long for CrossEntropyLoss

File /opt/conda/lib/python3.10/site-packages/tqdm/std.py:1181, in tqdm.__iter__(self)
   1178 time = self._time
   1180 try:
-> 1181     for obj in iterable:
   1182         yield obj
   1183         # Update and possibly print the progressbar.
   1184         # Note: does not call self.update(1) for speed optimisation.

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1344, in _MultiProcessingDataLoaderIter._next_data(self)
   1342 else:
   1343     del self._task_info[idx]
-> 1344     return self._process_data(data)

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1370, in _MultiProcessingDataLoaderIter._process_data(self, data)
   1368 self._try_put_index()
   1369 if isinstance(data, ExceptionWrapper):
-> 1370     data.reraise()
   1371 return data

File /opt/conda/lib/python3.10/site-packages/torch/_utils.py:706, in ExceptionWrapper.reraise(self)
    702 except TypeError:
    703     # If the exception takes multiple arguments, don't try to
    704     # instantiate since we don't know how to
    705     raise RuntimeError(msg) from None
--> 706 raise exception

ValueError: Caught ValueError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 50, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/dataset.py", line 420, in __getitems__
    return [self.dataset[self.indices[idx]] for idx in indices]
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/dataset.py", line 420, in <listcomp>
    return [self.dataset[self.indices[idx]] for idx in indices]
  File "/opt/conda/lib/python3.10/site-packages/torchvision/datasets/oxford_iiit_pet.py", line 116, in __getitem__
    image, target = self.transforms(image, target)
  File "/opt/conda/lib/python3.10/site-packages/torchvision/datasets/vision.py", line 97, in __call__
    target = self.target_transform(target)
  File "/opt/conda/lib/python3.10/site-packages/torchvision/transforms/transforms.py", line 95, in __call__
    img = t(img)
  File "/opt/conda/lib/python3.10/site-packages/torchvision/transforms/transforms.py", line 479, in __call__
    return self.lambd(img)
  File "/tmp/ipykernel_30/1474935402.py", line 39, in <lambda>
    transforms.Lambda(lambda mask: encode_segmap(mask)),
  File "/tmp/ipykernel_30/1474935402.py", line 22, in encode_segmap
    encoded_mask[np.all(mask == background, axis=-1)] = 0
ValueError: operands could not be broadcast together with shapes (256,256) (3,)