I have written a custom loss function and keep getting the error :
UnboundLocalError: local variable ‘beta1’ referenced before assignment
The error occurs when I use my loss function (see below) along with AdamW as my optimizer. According to the forums this was caused by a bug in PyTorch 1.8.0, so I upgraded to 1.8.1 but still the bug remains.
Below is the function which is throwing the error.
class SoftCel(nn.Module):
def __init__(self, ignore_val, class_weights = None, device = 'cpu'):
super(SoftCel, self).__init__()
self.device = device
self.ignore = ignore_val
self.class_weights = class_weights.to(self.device)
if self.class_weights is None:
self.class_weights = torch.ones((5, 1, 1)).to(self.device)
def get_class_weights_for_mean(self, labels):
"""If we use class weights we need to normalize by the weights
as per pytorch forum:
loss_weighted_manual = loss_weighted_manual.sum() / weights[target].sum()
Thus this function gets us the denominator of the above EQ.
"""
weights = []
classes = torch.unique(labels)
for c in classes:
if c == self.ignore:
continue
# number times a label is present
num = len(labels[labels == c])
# get weight value for that label
w = self.class_weights.squeeze()[c].item()
temp = [w for _ in range(num)]
weights.extend(temp)
return weights
def forward(self,
preds,
labels,
one_hot,
pixel_weights = None):
# list for pixel loss values and pixel weight values
L = []
W = []
for i in range(preds.shape[0]):
loss_val = - one_hot[i].long() * F.log_softmax(preds[i])
if pixel_weights is not None:
loss_val = loss_val * pixel_weights[i]
if self.class_weights is not None:
# #loss_val = self.class_weights * loss_val
loss_val = loss_val * self.class_weights
loss_val = torch.sum(loss_val, 0)
if self.ignore is not None:
loss_val = loss_val[labels[i]!=self.ignore]
L.extend(loss_val)
W.extend(self.get_class_weights_for_mean(labels[i]))
L = torch.tensor(L)
W = torch.tensor(W)
# get mean
loss = L.sum() / W.sum()
loss.requires_grad = True
return loss