[Question] Loss Function

[Soooda]

Hello buddy, I just started learning PointNet and am working on my own implementation currently as well. Your code has offered a great insight into my work but at the moment I don't quite understand your loss function setup somehow.

def pointnetloss(outputs, labels, m3x3, m64x64, alpha = 0.0001):
    criterion = torch.nn.NLLLoss()
    bs=outputs.size(0)
    id3x3 = torch.eye(3, requires_grad=True).repeat(bs,1,1)
    id64x64 = torch.eye(64, requires_grad=True).repeat(bs,1,1)
    if outputs.is_cuda:
        id3x3=id3x3.cuda()
        id64x64=id64x64.cuda()
    diff3x3 = id3x3-torch.bmm(m3x3,m3x3.transpose(1,2))
    diff64x64 = id64x64-torch.bmm(m64x64,m64x64.transpose(1,2))
    return criterion(outputs, labels) + alpha * (torch.norm(diff3x3)+torch.norm(diff64x64)) / float(bs)

According to the PointNet paper Page 4 Joint Alignment Network section and Supplementary Part C "A regularization loss (with weight 0.001) is added to the softmax classification loss to make the matrix close to orthogonal.", shouldn't the loss function be something like below:

# Regularization function from the official pytorch implementation repo 
# Ref. https://github.com/fxia22/pointnet.pytorch/blob/master/pointnet/model.py
def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))

def pointnetloss(outputs, labels, m3x3, m64x64, alpha=0.001):
    criterion = torch.nn.NLLLoss()
    return criterion(outputs, labels) + alpha * feature_transform_regularizer(m64x64)

I newly just stepped into this area of research and tbh I am still not confident to say that I fully understand the paper. If you can show me some pointers, that would be more than appreciated!
Cheers~