How to segment a part of images using neural network?

Hi, I’m trying to image segmentation using pytorch.

I have input images for dataset, and also ground truth images for label. So, my labels are just image files, not text.
And I don’t use any pre-trained dataset because I want to just train my own data. Most of examples in google used pre-trained dataset because they use text label for segmentation of objects, people, background, etc.

Anyway, I got error in CrossEntropyLoss function, and I would like to know how to get loss values between data and label .
I have no idea how to compare between input images and label images.

Here is my code. But it doesn’t work.

I got error -> RuntimeError(“the derivative for ‘target’ is not implemented”,)
in this line -> loss = criterion(output, output_target)

Could you help me to solve it or introduce me tutorial about image segmentation without pre-trained dataset using PyTorch?

class ImageDataset(Dataset):
    def __init__(self, is_train_set=False):
        trans_input = transforms.Compose([transforms.ToTensor()])
        trans_target = transforms.Compose([transforms.ToTensor()])
        trainset_input = torchvision.datasets.ImageFolder(root = './data/data01/image', transform=trans_input)
        trainset_target = torchvision.datasets.ImageFolder(root = './data/data01/label', transform=trans_target)
    
        len_train = len(trainset_input)
        trainloader_input = DataLoader(trainset_input, batch_size=len_train, shuffle=False)
        len_train = len(trainset_target)
        trainloader_target = DataLoader(trainset_target, batch_size=len_train, shuffle=False)

        dataiter_input = iter(trainloader_input)
        images_input, labels_input = dataiter_input.next()
        dataiter_target = iter(trainloader_target)
        images_target, labels_target = dataiter_target.next()
        #print('ITER;;; ', images_input[0])

        self.image = images_input
        self.target = images_target
        self.len = len_train

    def __getitem__(self, index):
        return self.image[index], self.target[index]

    def __len__(self):
        return self.len


# Training settings
batch_size = 64
train_dataset = ImageDataset(False)
train_loader = DataLoader(dataset=train_dataset,
                            batch_size=batch_size,
                            shuffle=True)

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.l1 = nn.Linear(100 * 100, 3600)
        self.l2 = nn.Linear(3600, 2500)
        self.l3 = nn.Linear(2500, 900)
        self.l4 = nn.Linear(900, 100)

    def forward(self, x):
        x = x.view(-1, 100 * 100)  # Flatten the data (n, 1, 28, 28)-> (n, 784)
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        return self.l4(x)


model = Net()
model_target = Net()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)


def train(epoch):
    model.train()
    model_target.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        output_target = model_target(target)
        loss = criterion(output, output_target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.data[0]))