Iterative loss function Autoencoders

Question

I am trying to implement a custom loss function in a Pytorch Autoencoder.

The loss function tries to maximize the cosine similarity between a given output tensor U (a vector) and 100 random vectors J where both U and J have the same dimension of [300]. This is repeated for each batch.

Suppose we have 30 items per batch, then the output tensor is

train_Y.shape = [30,300]

Random_vectors.shape = [30,100,300]

I can implement the loss function in two ways:

    All_Y =[]
    for Y,z_r in zip(train_y, random_vectors):
        Y_cosine_list =[]
        for z in z_r:
            cosi = torch.dot(Y,z) / (torch.norm(Y)*torch.norm(z))
            Y_cosine_list.append(cosi)
        All_Y.append(Y_cosine_list)
    
    All_Y = torch.tensor(All_Y).to(device)
    train_loss = torch.sum(torch.abs(All_Y))/dim_0
    train_loss = torch.tensor(train_loss.data, requires_grad = True)

or

    train_Y = torch.zeros([dim_0, 100])
    for i, (Y,z_r) in enumerate(zip(train_Y, random_vectors)):
        for j,z in enumerate(z_r):
            train_Y[i,j] = cos(Y,z)
    train_Y = train_Y.to(device)
    
    train_loss = torch.sum(torch.abs(train_Y))/dim_0

The second one is more elegant and to the point. However it is giving a "Cuda illegal memory access error". I have checked that the memory is not exceeded in either case. Is there anything wrong with the second implementation?

The first implementation is inelegant and I am not sure that it makes sense from a neural net optimization perspective. But it does not give errors and am able to complete training for all the epochs.

Ps: I have tried encapsulating this code block in a loss_fn method but I get the same illegal memory access error.

I have tried everything that I could find for the illegal memory access error - changing GPUs, removing a torch.stack block etc. But I can't seem to get rid of the problem.

user_14108734 · Answer

Here is a vectorized way to do it

class CosineLoss(nn.Module):
    def __init__(self, ):
        super().__init__()
        pass

    def forward(self, x, y):
        """
        Args:
            x (torch.tensor): [batchsize, N, M] - tensor.
            y (torch.tensor): [batchsize, M] - tensor.

        Returns:
            torch.tensor: scalar mean cosine loss
        """
        # dot product along dimension 'm' i.e multiply and sum along 'm'.
        dotp = torch.einsum("bm, bnm -> bn", y, x)
        # L2 norm along dimension 'm' and multiply by broadcasting
        length = torch.norm(y, dim=-1)[:, None]*torch.norm(x, dim=-1)
        # cosine = dotproduct of unit vectors
        cos = dotp/length
        return cos.mean()


def test():
    b, n, m = 30, 100, 300
    train_Y = torch.randn(b, m, device='cuda')
    random_vectors = torch.randn(b, n, m, requires_grad=True, device='cuda')
    print(f'{random_vectors.grad = }')

    cosineloss = CosineLoss()
    loss = cosineloss(random_vectors, train_Y)
    print(f'{loss = }')

    loss.backward()
    print(f'{random_vectors.grad.shape = }')

References:

Iterative loss function Autoencoders

1 Answers