Custom weighted loss function

Hello everyone, I’m having a “weighted” problem with this particular loss function.

Context:
I’m training a capsule network model with two classes, but the classes are unbalanced(class 0 has 32% of the training samples of the dataset and class 1 has 68%).

I’m using inversed weights to force the network to learn from the class 0 data, but the model gets 68% of accuracy all the time, it’s just learning from the class 1 (just in the capsule network model).

I have tried with an InceptionV3 with weighted cross entropy loss as criterion, to see if the weighted works:

criterion = nn.CrossEntropyLoss(torch.FloatTensor([0.68, 0.32]).cuda())

And the model gets 98% accuracy, it works.

But when I try to weight the capsule network loss function like this:

def margin_loss(self, x, labels, size_average=True):

	batch_size = x.size(0)
	
	v_c = torch.sqrt((x**2).sum(dim=2, keepdim=True))
	
	left = F.relu(0.9 - v_c).view(batch_size, -1)
	
	right = F.relu(v_c - 0.1).view(batch_size, -1)

	loss = labels * left + 0.5 * (1.0 - labels) * right

	# This line is not included in the original code, it's the "weighted" that i want to apply
	loss = loss * self.weights
	
	loss = loss.sum(dim=1).mean()

	return loss

Nothing happen, is still the same accuracy.

I have tried with the following weights:

self.weights = torch.Tensor([0.68, 0.32]).cuda() # 68% accuracy in all epochs
self.weights = torch.Tensor([0.99, 0.01]).cuda() # Check the attached image

But the model is still stuck on 68% accuracy(class 1).

Look at this image, on the right are the inceptionv3 results, on the left the capsule network results in [loss, accuracy] format, every record is an epoch:

The code is based on this capsule networks repo

You can check the complete model in there.

Any idea how to weight that particular loss function? I will be reading your answers! Thanks!