# Importing Essential Libraries
import os
import numpy as np
import time
import copy
import random
​
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.datasets import OxfordIIITPet
import torchvision.transforms.functional as TF
​
import matplotlib.pyplot as plt
from PIL import Image
​
# Check GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
Using device: cuda
add Codeadd Markdown
# Custom Dataset Class
class PetDataset(Dataset):
    def __init__(self, root, transform=None, target_transform=None):
        self.dataset = OxfordIIITPet(root=root, download=True, target_types='segmentation')
        self.transform = transform
        self.target_transform = target_transform
​
    def __len__(self):
        return len(self.dataset)
​
    def __getitem__(self, idx):
        # Unpack only two elements: image and mask
        image, mask = self.dataset[idx]
        
        # Convert mask to binary: pet vs. background
        mask = np.array(mask)
        # In Oxford-IIIT Pet Dataset:
        # Pixel values:
        #   1 = foreground (pet)
        #   2 = background
        #   3 = boundary
        # We'll convert all non-pet pixels to background for binary segmentation
        mask = np.where(mask == 1, 1, 0).astype(np.float32)
        mask = Image.fromarray(mask)
​
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            mask = self.target_transform(mask)
        
        return image, mask
add Codeadd Markdown
# Define Image and Mask Transformations
image_size = 256
​
# Transformations for training data
train_transforms = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.RandomResizedCrop(image_size, scale=(0.8, 1.0)),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],  # ImageNet mean
                         [0.229, 0.224, 0.225])  # ImageNet std
])
​
# Transformations for validation data
val_transforms = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])
​
# Transformations for masks
mask_transforms = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.ToTensor()
])
add Codeadd Markdown
# Initialize the Dataset
dataset = PetDataset(root='./data', transform=train_transforms, target_transform=mask_transforms)
​
# Split the Dataset into Training and Validation Sets
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
​
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
​
# Update Transforms for Validation Dataset (No Augmentations)
val_dataset.dataset.transform = val_transforms
​
# Define DataLoaders with Reduced Batch Size
batch_size = 2  # Reduced from 8 to 2 to save memory
​
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)
​
print(f"Number of training samples: {len(train_dataset)}")
print(f"Number of validation samples: {len(val_dataset)}")
Number of training samples: 2944
Number of validation samples: 736
add Codeadd Markdown
# CBAM (Convolutional Block Attention Module)
class CBAM(nn.Module):
    def __init__(self, channels, reduction=16, kernel_size=7):
        super(CBAM, self).__init__()
        # Channel Attention Module
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, channels // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channels // reduction, channels, 1, bias=False),
            nn.Sigmoid()
        )
        # Spatial Attention Module
        self.spatial_attention = nn.Sequential(
            nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size//2, bias=False),
            nn.Sigmoid()
        )
​
    def forward(self, x):
        # Channel Attention
        ca = self.channel_attention(x)
        x = x * ca
        # Spatial Attention
        max_pool, avg_pool = torch.max(x, dim=1, keepdim=True)[0], torch.mean(x, dim=1, keepdim=True)
        sa = self.spatial_attention(torch.cat([max_pool, avg_pool], dim=1))
        x = x * sa
        return x
​
# Self-Attention Module
class SelfAttention(nn.Module):
    def __init__(self, in_dim):
        super(SelfAttention, self).__init__()
        self.channel_in = in_dim
​
        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)
        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)
        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
​
        self.gamma = nn.Parameter(torch.zeros(1))
​
        self.softmax = nn.Softmax(dim=-1)
​
    def forward(self, x):
        """
        inputs:
            x: input feature maps (B X C X W X H)
        returns:
            out: self attention value + input feature
            attention: B X N X N (N = W*H)
        """
        m_batchsize, C, width, height = x.size()
        proj_query = self.query_conv(x).view(m_batchsize, -1, width*height).permute(0, 2, 1)  # B X (W*H) X C'
        proj_key = self.key_conv(x).view(m_batchsize, -1, width*height)  # B X C' X (W*H)
        energy = torch.bmm(proj_query, proj_key)  # B X (W*H) X (W*H)
        attention = self.softmax(energy)  # B X (W*H) X (W*H)
        proj_value = self.value_conv(x).view(m_batchsize, -1, width*height)  # B X C X (W*H)
​
        out = torch.bmm(proj_value, attention.permute(0, 2, 1))  # B X C X (W*H)
        out = out.view(m_batchsize, C, width, height)
​
        out = self.gamma * out + x
        return out
add Codeadd Markdown
# Double Convolution Block with Dropout
class DoubleConv(nn.Module):
    """(Conv2d => BatchNorm => ReLU) * 2 with optional Dropout"""
​
    def __init__(self, in_channels, out_channels, dropout_prob=0.1):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout_prob),
​
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
​
    def forward(self, x):
        return self.double_conv(x)
​
# U-Net Module
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, features=[32, 64, 128]):
        super(UNet, self).__init__()
        self.downs = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
​
        # Down part of UNet
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))
            in_channels = feature
​
        # Bottleneck
        self.bottleneck = DoubleConv(features[-1], features[-1]*2)
​
        # Up part of UNet
        for feature in reversed(features):
            self.ups.append(
                nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2)
            )
            self.ups.append(DoubleConv(feature*2, feature))
​
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
​
    def forward(self, x):
        skip_connections = []
​
        for down in self.downs:
            x = down(x)
            skip_connections.append(x)
            x = self.pool(x)
​
        x = self.bottleneck(x)
​
        skip_connections = skip_connections[::-1]
​
        for idx in range(0, len(self.ups), 2):
            x = self.ups[idx](x)
            skip_connection = skip_connections[idx//2]
​
            if x.shape != skip_connection.shape:
                x = TF.resize(x, size=skip_connection.shape[2:])
​
            concat_skip = torch.cat((skip_connection, x), dim=1)
            x = self.ups[idx+1](concat_skip)
​
        return self.final_conv(x)
​
# U²-Net Module
class U2Net(nn.Module):
    def __init__(self, in_channels=3, out_channels=512, features=[32, 64, 128]):
        super(U2Net, self).__init__()
        self.encoder = UNet(in_channels, out_channels, features=features)
        self.decoder = UNet(out_channels, out_channels, features=features)
        self.attention = CBAM(out_channels)  # CBAM expects 'out_channels' channels
​
    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        x = self.attention(x)
        return x
add Codeadd Markdown
# U3Net Module
class U3Net(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, features=[32, 64, 128], dropout_prob=0.1):
        super(U3Net, self).__init__()
​
        # First Hierarchical Level with CBAM Attention
        self.u2net1 = U2Net(in_channels, out_channels=512, features=features)  # Output: 512 channels
        self.unet1 = UNet(512, 512, features=features)  # Output: 512 channels
        self.cbam1 = CBAM(512)  # CBAM expects 512 channels
​
        # Second Hierarchical Level with Self-Attention
        self.u2net2 = U2Net(512, out_channels=512, features=features)  # Output: 512 channels
        self.unet2 = UNet(512, 512, features=features)  # Output: 512 channels
        self.self_attention2 = SelfAttention(512)  # Self-Attention expects 512 channels
​
        # Final Convolution for Output
        # Concatenated channels = 512 + 512 = 1024
        self.final_conv = nn.Conv2d(512*2, out_channels, kernel_size=1)
​
    def forward(self, x):
        # First Level Processing
        x1 = self.u2net1(x)  # 512 channels
        x1 = self.unet1(x1)   # 512 channels
        x1 = self.cbam1(x1)   # 512 channels
​
        # Second Level Processing
        x2 = self.u2net2(x1)  # 512 channels
        x2 = self.unet2(x2)    # 512 channels
        x2 = self.self_attention2(x2)  # 512 channels
​
        # Combine Features
        combined = torch.cat([x1, x2], dim=1)  # Concatenated to 1024 channels
​
        # Final Prediction
        out = self.final_conv(combined)  # 1 channel
        out = torch.sigmoid(out)  # Binary segmentation
​
        return out
add Codeadd Markdown
# Dice Coefficient
def dice_coefficient(pred, target, smooth=1.):
    pred = pred.contiguous()
    target = target.contiguous()
    
    intersection = (pred * target).sum(dim=2).sum(dim=2)
    dice = (2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)
    
    return dice.mean()
​
# Dice Loss
class DiceLoss(nn.Module):
    def __init__(self):
        super(DiceLoss, self).__init__()
    
    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        pred = pred.view(pred.size(0), -1)
        target = target.view(target.size(0), -1)
        
        intersection = (pred * target).sum(1)
        dice = (2. * intersection + 1) / (pred.sum(1) + target.sum(1) + 1)
        loss = 1 - dice
        return loss.mean()
add Codeadd Markdown
# Checkpointing Utilities
def save_checkpoint(state, filename="checkpoint.pth.tar"):
    torch.save(state, filename)
​
def load_checkpoint(checkpoint, model, optimizer):
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    epoch = checkpoint['epoch']
    best_loss = checkpoint['best_loss']
    return epoch, best_loss
add Codeadd Markdown
# Early Stopping Implementation
class EarlyStopping:
    def __init__(self, patience=10, verbose=False, delta=0):
        """
        Args:
            patience (int): How long to wait after last improvement.
            verbose (bool): If True, prints a message for each validation loss improvement.
            delta (float): Minimum change 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
​
    def __call__(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
            if self.verbose:
                print(f"Initial validation loss: {self.best_loss:.4f}")
            return
​
        elif val_loss < self.best_loss - self.delta:
            self.best_loss = val_loss
            self.counter = 0
            if self.verbose:
                print(f'Validation loss decreased to {self.best_loss:.4f}. Resetting counter.')
        else:
            self.counter += 1
            if self.verbose:
                print(f'No improvement in validation loss for {self.counter} epochs.')
            if self.counter >= self.patience:
                self.early_stop = True
add Codeadd Markdown
from torch.amp import autocast, GradScaler  # Updated imports for mixed precision
​
def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs=50, patience=10, checkpoint_path='checkpoint.pth.tar'):
    start_epoch = 0
    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = float('inf')
    
    # Initialize Early Stopping
    early_stopping = EarlyStopping(patience=patience, verbose=True)
    
    # Initialize GradScaler for Mixed Precision Training
    scaler = GradScaler()
    
    # Check if checkpoint exists
    if os.path.exists(checkpoint_path):
        print("Loading checkpoint...")
        checkpoint = torch.load(checkpoint_path, map_location=device)
        start_epoch, best_loss_loaded = load_checkpoint(checkpoint, model, optimizer)
        best_loss = best_loss_loaded
        print(f"Resuming from epoch {start_epoch+1}")
    
    for epoch in range(start_epoch, num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
                dataloader = train_loader
            else:
                model.eval()
                dataloader = val_loader
            
            running_loss = 0.0
            running_dice = 0.0
            
            # Iterate over data
            for inputs, masks in dataloader:
                inputs = inputs.to(device, non_blocking=True)
                masks = masks.to(device, non_blocking=True)
                
                # Zero the parameter gradients
                optimizer.zero_grad()
                
                # Forward with mixed precision
                with autocast():
                    outputs = model(inputs)
                    loss = criterion(outputs, masks)
                    dice = dice_coefficient(outputs, masks)
                
                # Backward + optimize only if in training phase
                if phase == 'train':
                    scaler.scale(loss).backward()
                    # Gradient Clipping
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                    scaler.step(optimizer)
                    scaler.update()
                
                running_loss += loss.item() * inputs.size(0)
                running_dice += dice.item() * inputs.size(0)
            
            epoch_loss = running_loss / len(dataloader.dataset)
            epoch_dice = running_dice / len(dataloader.dataset)
            
            print(f'{phase.capitalize()} Loss: {epoch_loss:.4f} Dice: {epoch_dice:.4f}')
            
            # Deep copy the model
            if phase == 'val':
                scheduler.step(epoch_loss)
                if epoch_loss < best_loss:
                    best_loss = epoch_loss
                    best_model_wts = copy.deepcopy(model.state_dict())
                # Early Stopping
                early_stopping(epoch_loss)
                save_checkpoint({
                    'epoch': epoch,
                    'state_dict': model.state_dict(),
                    'optimizer' : optimizer.state_dict(),
                    'best_loss': best_loss
                }, filename=checkpoint_path)
                
                if early_stopping.early_stop:
                    print("Early stopping triggered!")
                    model.load_state_dict(best_model_wts)
                    return model
        
        print()
    
    # Load best model weights
    model.load_state_dict(best_model_wts)
    return model
add Codeadd Markdown
# Initialize the U3Net Model with Reduced Features
model = U3Net(in_channels=3, out_channels=1, features=[32, 64, 128]).to(device)
​
# Define Loss Function and Optimizer
criterion = DiceLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)  # Added weight_decay for regularization
​
# Define Learning Rate Scheduler
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
​
# Define Training Parameters
num_epochs = 50
patience = 10
checkpoint_path = 'u3net_checkpoint.pth.tar'
​
# Start Training
trained_model = train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, 
                            num_epochs=num_epochs, patience=patience, checkpoint_path=checkpoint_path)
​
# Save the Best Model Weights
torch.save(trained_model.state_dict(), 'u3net_best_model.pth')
print("Training complete. Best model saved as 'u3net_best_model.pth'")
Epoch 1/50
----------
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], line 17
     14 checkpoint_path = 'u3net_checkpoint.pth.tar'
     16 # Start Training
---> 17 trained_model = train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, 
     18                             num_epochs=num_epochs, patience=patience, checkpoint_path=checkpoint_path)
     20 # Save the Best Model Weights
     21 torch.save(trained_model.state_dict(), 'u3net_best_model.pth')

Cell In[15], line 47, in train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, patience, checkpoint_path)
     44 optimizer.zero_grad()
     46 # Forward with mixed precision
---> 47 with autocast():
     48     outputs = model(inputs)
     49     loss = criterion(outputs, masks)

TypeError: autocast.__init__() missing 1 required positional argument: 'device_type'
add Codeadd Markdown
arrow_upwardarrow_downwarddelete
# Visualization Function
def visualize_predictions(model, dataloader, num_samples=5):
    model.eval()
    samples = 0
    plt.figure(figsize=(15, num_samples * 5))
    
    with torch.no_grad():
        for inputs, masks in dataloader:
            inputs = inputs.to(device)
            masks = masks.to(device)
            outputs = model(inputs)
            preds = (outputs > 0.5).float()
            
            for i in range(inputs.size(0)):
                if samples >= num_samples:
                    return
                img = inputs[i].cpu().numpy().transpose(1, 2, 0)
                # Revert normalization
                img = np.clip(img * np.array([0.229, 0.224, 0.225]) + 
                              np.array([0.485, 0.456, 0.406]), 0, 1)
                
                mask = masks[i].cpu().numpy().squeeze()
                pred = preds[i].cpu().numpy().squeeze()
                
                plt.subplot(num_samples, 3, samples*3 + 1)
                plt.imshow(img)
                plt.title("Input Image")
                plt.axis('off')
                
                plt.subplot(num_samples, 3, samples*3 + 2)
                plt.imshow(mask, cmap='gray')
                plt.title("Ground Truth")
                plt.axis('off')
                
                plt.subplot(num_samples, 3, samples*3 + 3)
                plt.imshow(pred, cmap='gray')
                plt.title("Prediction")
                plt.axis('off')
                
                samples += 1
​
# Visualize Predictions on the Validation Set
visualize_predictions(trained_model, val_loader, num_samples=5)
plt.tight_layout()
plt.show()


help me with the issue