torch.argmax() removes grad_fn 02

[mo3az-14]

Code

video no .88 Time stamp : 13:27:57

My error

I keep getting errors in the loss.backward() saying element 0 of tensors does not require grad and does not have a grad_fn after some research and troubleshooting I found out that torch.argmax( ) removes the grad_fn on the tensor thus it becomes unable to backpropagation

My code

X_train , X_test = X_train.to(device) , X_test.to(device)
y_train, y_test , y_train.to(device),y_test.to(device)
epochs = 100
#training loop
for epoch in range(epochs)  : 
  #training
  model_0.train()
  #forward pass
  y_logits = model_0(X_train)
  y_preds = torch.softmax(y_logits, dim=1)
  y_preds = y_preds.argmax(dim=1).type(torch.float) # <<----- WHY DOES THE grad_fn DISAPPEAR
  #print(y_preds.grad_fn)  output: None  
  #calculate loss
  loss = loss_fn( y_preds ,y_train )
  #calculate accuracy
  #acc =accuracy_fn(y_true = y_train , y_pred = y_preds ) <<-- i keep getting an error here but not that much of a deal so i just commented it out 
  #zero grad
  optimizer.zero_grad()
  #back propagation
  loss.backward()
  #optimizer step
  optimizer.step()
  #testing phase
  model_0.eval()
  with torch.inference_mode(): 
    test_logits = model_0(X_test.type(torch.float))
    test_preds= torch.softmax(test_logits,dim=1).argmax(dim=1).type(torch.float) #<<--- grad_fn doesn't matter here since it's just testing
    test_acc = accuracy_fn(y_pred= test_preds ,y_true= y_test)
    test_loss = loss_fn(test_preds , y_test.type(torch.float))
  #print results  
  if epoch%10 == 0 : 
    print(f"EPOCHS:{epoch} | train loss:{loss:.4f} | train accuracy: N/A % | test loss:{test_loss :.4f} | test accuracy:{test_acc:.2f}%")

stuff I tried

1.I tried using torch.requires_grad_() forcing it gradient accumlation but it didn't work since it starts a new gradient
2.I tried using torch.max( ) and doing a for loop (victorized) but it also removes the grad_fn

[nazarPuriy]

Some theorical explanation
I think I have the answer. Let my try to explain. When you do backward() you are calculating gradients and for that you need a differentiable function, usual type of functions: x², sin(x) etc. Backward is used when you have a function that goes from R^n to R. For example you put (1,2,3) and it outputs: 0.56. Then you can imagine the function living in its space and you want to find the minimum or maximum, with bacwards you obtain the gradient in each tensor. You can see the gradient as a vector that points in which direction the slope is higher.; then with step you move in that direction. For all that you need smooth function (differentiable functions). https://www.geogebra.org/3d/skqjgzba here you can see a smooth and b not differentiable. So when the function is not differentiable it can happen that you cannot obtain the gradient and then cannot do the step and so on.

**Practical
argmax is not differentiable-> you can't calculate gradients on it -> you can't do backward on a function that has it.
since your loss depends on y_preds that depends on argmax your loss is not differentiable and cannot do backwards to obtain gradients.
Your lose should be
loss = loss_fn( y_logits ,y_train ) // Before the softmax or after depending the loss function (as some of them already incorporate the softmax; probably you don't need softmax).

[mo3az-14]

Thank you for the explanation! ,But the problem is why did it work for mrdbourke😅. I tried copying his code bit by bit to understand where did it all go wrong but still couldn't figure it out. I saw some people suggesting to just ignore the .argmax() step since it is not critical and only apply it when calculating accuracy, in that case we ignore the loss as I would get 220.4113 loss and 100% accuracy, and that's what I did.

[mrdbourke]

Hi @mo3az-14,

Wha happens if you use some_tensor.argmax() instead of torch.argmax(), have you tried that?

For example:

for epoch in range(epochs):
    ### Training
    model_4.train()

    # 1. Forward pass
    y_logits = model_4(X_blob_train) # model outputs raw logits 
    y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1) # go from logits -> prediction probabilities -> prediction labels
    # print(y_logits)
    # 2. Calculate loss and accuracy
    loss = loss_fn(y_logits, y_blob_train) 
    acc = accuracy_fn(y_true=y_blob_train,
                      y_pred=y_pred)

I just ran through all of the code in notebook 02 and it functions as expected: https://github.com/mrdbourke/pytorch-deep-learning/blob/main/02_pytorch_classification.ipynb

Update (from @nazarPuriy):

Calculate the loss on the y_logits rather than y_preds.

Good:

# Calculate loss on logits
loss = loss_fn(y_logits, y_train)

Bad:

# Calculate loss on preds
loss = loss_fn(y_preds, y_train)

[mo3az-14]

yes!, I have already tried it but same issue . when I copied your code it work , but for some reason mine didn't work. I tried copying your code bit by bit to try and see why is it not working but unfortunately it didn't work.
thank you tho <3.

[nazarPuriy]

Because you are performing the loss to the function that incorporates argmax:

His code:

    y_logits = model_4(X_blob_train) # model outputs raw logits 

    y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1) 

    loss = loss_fn(y_logits, y_blob_train) 

The loss is done with y_logits that don't have the argmax

Your code:

  y_logits = model_0(X_train)
  y_preds = torch.softmax(y_logits, dim=1)
  y_preds = y_preds.argmax(dim=1).type(torch.float) 
  loss = loss_fn( y_preds ,y_train )

You are doing the loss with y_preds that does have the argmax
Just change
loss = loss_fn( y_preds ,y_train )
by
loss = loss_fn( y_logits ,y_train )

It's not because of the usage of the argmax, is because the calculus of the loss and later the gradients using the argmax inside the loss function.

Hope it works!

[mo3az-14]

It worked , I didn't know i was supposed to use the raw logits not the prediction s , thanks alot !

[mrdbourke]

@nazarPuriy thank you for the update!