Re-training deeplab v3

vision
1

Can someone help me with a link to a tutorial on how to re-training deeplab v3 on my data? I have only one class target and I keep getting errors. I’m fairly new to pytorch. I have been searching and reading but still unsucessful. TIA! Currently my code is at this stage:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.transforms import ToTensor
from torch.utils.data import Dataset, DataLoader
from torchvision.models.segmentation import deeplabv3_resnet50
import os
from PIL import Image
from torchvision import transforms

# Define your own dataset
class MyDataset(Dataset):
    def __init__(self, image_paths, mask_paths, transform=None):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform

    def __getitem__(self, index):
        image = Image.open(self.image_paths[index]).convert("RGB")
        mask = Image.open(self.mask_paths[index])
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)
        return image, mask

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


# Define your training hyperparameters
batch_size = 16
learning_rate = 0.001
num_epochs = 10

# Create an instance of the DeepLabV3 model
model = deeplabv3_resnet50(pretrained=False, num_classes=2)  # Adjust num_classes to 2 for object and background

# Define your loss function (e.g., CrossEntropyLoss) and optimizer (e.g., Adam)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)


myTransform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224, 224))
    #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


# Load your dataset
image_dir = "./test_img_mask/images"
mask_dir = "./test_img_mask/masks"
image_paths = [os.path.join(image_dir, filename) for filename in os.listdir(image_dir)]
mask_paths = [os.path.join(mask_dir, filename) for filename in os.listdir(mask_dir)]
dataset = MyDataset(image_paths, mask_paths, transform=myTransform)

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])

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)

# Move the model to the GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Training loop
for epoch in range(num_epochs):
    model.train()
    for images, masks in train_loader:
        images = images.to(device)
        masks = masks.to(device)

        # Forward pass
        outputs = model(images)['out']

        # Compute loss
        loss = criterion(outputs, masks.squeeze(1).long())

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    # Validation loop
    model.eval()
    with torch.no_grad():
        for images, masks in val_loader:
            images = images.to(device)
            masks = masks.to(device)

            # Forward pass
            outputs = model(images)['out']

            # Compute validation loss
            val_loss = criterion(outputs, masks.squeeze(1).long())

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}, Val Loss: {val_loss.item():.4f}')