How should I implement my very own custom loss function.
There are two ways to implement using a custom loss class and writing function with init
and forward
.
I had the code already written for TensorFlow in class, therefore I chose to go along.
For PyTorch, I tried explored the below relevant discussions, but could not figure how to design the backwards function for same.
Custom Loss:
class MSE_ITL_RegLoss(torch.nn.Module):
def __init__(self, latent_code, prior, regularization_factor, kernel, batch_size):
super (MSE_ITL_RegLoss).__init__()
self.latent_code = latent_code
self.regularization_factor = regularization_factor
self.prior = prior
self.kernel = kernel
self.batch_size = batch_size
def evaluate_ITL(self, first_pdf, second_pdf):
error = first_pdf - second_pdf
pairwise_error = torch.reshape(torch.unsqueeze(error ,1) - error, (-1,))
temp = torch.exp(-(np.sqrt(2) * self.kernel) * torch.square(pairwise_error_t))
return torch.sum(temp)/self.batch_size**2
def cs_div (self):
p_x = self.evaluate_ITL(self.latent_code, self.latent_code)
q_x = self.evaluate_ITL(self.prior,self.prior)
p_x__q_x = self.evaluate_ITL(self.latent_code, self.prior)
return p_x__q_x / torch.sqrt(p_x * q_x)
def forward(self, y_true, y_pred):
mse = torch.mean(torch.square(y_true_t-y_pred_t))
cs_div = self.cs_div()
final_loss = mse + self.regularization_factor * cs_div
return final_loss
Model:
class modelMSE_CSD (nn.Module):
def __init__(self, input_size, bottleneck_dim):
super(modelMSE_CSD, self).__init__()
self.bottleneck_dim = bottleneck_dim
self.input_size = input_size
self.encoder = nn.Sequential(
nn.Linear(self.input_size, 800),
nn.ReLU(),
nn.Linear(800,400),
nn.ReLU(),
nn.Linear(400, self.bottleneck_dim)
)
self.decoder = nn.Sequential(
nn.Linear(self.bottleneck_dim, 400),
nn.ReLU(),
nn.Linear(400,800),
nn.ReLU(),
nn.Linear(800, self.input_size),
nn.Sigmoid()
)
def forward(self,x):
e_out = self.encoder(x)
d_out = self.decoder(e_out)
return e_out, d_out
Training Loop:
losses = []
for epoch in range(EPOCHS):
running_loss = 0.0
for batch_num, (data, target) in enumerate(dataloader):
data = Variable(data, requires_grad = True)
target = Variable(target, requires_grad = True)
model_csd.train()
latent_code, decoded_layer = model_csd(data)
loss = MSE_ITL_RegLoss(latent_code=latent_code, prior=PRIOR_DIST,
regularization_factor=REGULARIZER, kernel=KERNEL, batch_size=BATCH_SIZE)
optimizer.zero_grad()
**loss.backward()**
optimizer.step()
running_loss += loss.item() * data.size[0]
losses.append(running_loss/len(dataset))
print("Epoch: ",epoch,"Loss :", running_loss/len(dataset))```
Please suggest if the other method of declaring the custom loss is more easier.
Any links/resources would be appreciated.