Hello,
I am trying to train a semantic segmentation network for the first time and in general I am still quite new with PyTorch.
My inputs are RGB-images and corresponding grayscale images. In them the grayscale intensity corresponds to the pixel class.
Please see my first try in the following.
Right now I am getting the error: “Expected 4-dimensional input for 4-dimensional weight 64 3 7 7, but got 3-dimensional input of size [1, 1000, 1000] instead” from the line “output = model(img)” in the training loop.
The input with 1x1000x100 is the mask with the class labels.
How can this be solved and do you have further advises for me?
Thank you very much 
import torch
import os
import numpy as np
import torch.nn as nn
import torch.optim as optim
# DataLoader
from PIL import Image
from torchvision import transforms
from torch.utils.data.sampler import SubsetRandomSampler
class Dataset(torch.utils.data.Dataset):
def __init__(self, root, transforms=None):
self.root = root
self.transforms = transforms
# load all image files, sorting them to ensure that they are aligned
self.imgs = list(sorted(os.listdir(os.path.join(root, "Images")))) # Subfolder with images
self.masks = list(sorted(os.listdir(os.path.join(root, "ImageLabels")))) # Subfolder with masks
def __getitem__(self, idx):
# load images ad masks
img_path = os.path.join(self.root, "Images", self.imgs[idx])
mask_path = os.path.join(self.root, "ImageLabels", self.masks[idx])
img = Image.open(img_path).convert("RGB") # Three channel rgb images
mask = Image.open(mask_path) # Single channel and the pixelintensity corresponds to the class label
if self.transforms is not None:
img = self.transforms(img)
mask = self.transforms(mask)
img =transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)
# No need to normalize masks, they are all equal
return img, mask
def __len__(self):
return len(self.imgs)
transformations = transforms.Compose([
transforms.CenterCrop(1000),
transforms.ToTensor()
])
# Load data
dataset = Dataset('/images/', transforms=transformations)
# Split data (Random)
train_size = int(0.8 * len(dataset))
test_size = len(dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
dataloader_train=torch.utils.data.DataLoader(train_dataset,
batch_size=10, shuffle=True)
dataloader_test=torch.utils.data.DataLoader(test_dataset,
batch_size=10, shuffle=False)
# Model:
from torchvision import models
model = models.segmentation.deeplabv3_resnet101(pretrained=False, progress=True, num_classes=output_classes)
criterion = nn.BCEWithLogitsLoss().cuda
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model.train()
# Use gpu if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model.to(device)
# Training
num_epochs = 1
for iter in range(num_epochs):
# Training:
model.train()
for i in range(len(train_loader)):
img, mask = train_loader[i]
img, mask = img.cuda(), mask.cuda()
optimizer.zero_grad()
img = img.to(device)
mask = mask.to(device)
output = model(img)
loss = criterion(output, mask)
print(loss.cpu())
loss.backward()
optimizer.step()
# Validation:
model.eval()
for i in range(len(test_loader)):
img, mask = test_loader[i]
img, mask = img.cuda(), mask.cuda()
output = model.forward(img) # forward pass
loss = criterion(output, mask)
print(loss.cpu())
