I’m attempting to train a NN, but am getting an error when I attempt to train it. The data I’m using is .tif images and masks. The images are 3 band. I’ve tried to follow other questions which encountered the same issue to no avail. Any help would be greatly appreciated! I receive the following error when running:

ValueError                                Traceback (most recent call last)
<ipython-input-81-60713cba9810> in <cell line: 111>()
    123 
    124     # Calculate loss
--> 125     loss = criterion(outputs, masks)
    126 
    127     # Backward pass and optimization

3 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3084     if size_average is not None or reduce is not None:
   3085         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3086     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3087 
   3088 

ValueError: Expected input batch_size (4) to match target batch_size (262144).

# Define your model architecture
class UNet(nn.Module):
    def __init__(self, num_classes):
        super(UNet, self).__init__()
        # Define the architecture here
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1)
        self.fc = nn.Linear(64 * 256 * 256, num_classes)  # Adjust the output size

    def forward(self, x):
        # Define the forward pass here
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(x.size(0), -1)  # Flatten the output for fully connected layer
        x = self.fc(x)
        return x

# Define custom dataset class
from torchvision.transforms.functional import to_tensor

class CustomDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform

        # Get the list of image and mask files
        self.image_files = os.listdir(image_dir)
        self.mask_files = os.listdir(mask_dir)

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

    def __getitem__(self, idx):
        # Load image and mask
        image_path = os.path.join(self.image_dir, self.image_files[idx])
        mask_path = os.path.join(self.mask_dir, self.mask_files[idx])

        with rasterio.open(image_path) as image_src, rasterio.open(mask_path) as mask_src:
            image = image_src.read().astype(np.float32)
            mask = mask_src.read().astype(np.uint8)

        # Scale the image values to [0, 255] range
        image = (image - image.min()) / (image.max() - image.min()) * 255
        image = image.astype(np.uint8)

        # Convert NumPy arrays to PIL images
        image = Image.fromarray(np.moveaxis(image, 0, -1))
        mask = Image.fromarray(mask.squeeze())

        # Apply transformations if specified
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)

        # Convert PIL Images to tensors
        image = to_tensor(image)
        mask = to_tensor(mask)

        return image, mask



# Define transformations
transform = transforms.Compose([
    Resize((256, 256)),  # Resize images and masks to (256, 256)
    # Add any other transformations here
])

# Set file paths
image_dir = "/content/patches256/images/"
mask_dir = "/content/patches256/mask/"

# Split data into training and testing sets
train_image_files, test_image_files, train_mask_files, test_mask_files = train_test_split(
    os.listdir(image_dir), os.listdir(mask_dir), test_size=0.2, random_state=42)

# Create datasets for training and testing
train_dataset = CustomDataset(image_dir=image_dir, mask_dir=mask_dir, transform=transform)
test_dataset = CustomDataset(image_dir=image_dir, mask_dir=mask_dir, transform=transform)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=4, shuffle=False)

# Initialize the model with the correct number of classes
# Load a sample mask (assuming it's loaded as a NumPy array)
mask = np.array(...)  # Load your mask here

# Count the unique values
num_classes = len(np.unique(mask))
model = UNet(num_classes)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
for images, masks in train_loader:
    # Zero the parameter gradients
    optimizer.zero_grad()

    # Forward pass
    outputs = model(images)

    # Ensure masks are in the correct format
    masks = masks.squeeze(1).long().view(-1)  # Squeeze the channel dimension, convert to long tensor, and flatten

    # Resize outputs to match the shape of masks
    outputs = outputs.contiguous().view(-1, num_classes)

    # Calculate loss
    loss = criterion(outputs, masks)

    # Backward pass and optimization
    loss.backward()
    optimizer.step()

    # Print statistics
    running_loss += loss.item()




# Evaluate on test set
model.eval()
test_loss = 0.0
for images, masks in test_loader:
    outputs = model(images)
    masks = masks.squeeze(1).long()
    test_loss += criterion(outputs, masks).item()

print(f"Test Loss: {test_loss/len(test_loader)}")