A problem about the code

[draym28]

Hello, I read your code, and it seems that you forgot to self.zero_grad() before self.criterion(self.model(samples), targe). Is that a bug? Or I just miss something?

[rollovd]

Hey, man!

I'm sorry for a late response.
No, actually, there is no bug in the code.

Typically:

1. zero_grad
2. backward
3. step

My implementation:

1. backward
2. step
3. zero_grad

In general, this is the same, but in my implementation zero_grad will be called on the next step.
By "typical" implementation zero_grad will be redundant in the first step (and then it will work like mine).

So, my point is that both these approaches are cycled :)

[draym28]

Sorry for the late reply. :)
Yes, I understand that putting the zero_grad() in the end is also cycled.
I realized that I didn't clear my question.

The normal SAM procedure is:

• 1st backward
• compute e_w
• zero_grad
• 2nd backward
• update the original params
• zero_grad

So my point is that there should be a zero_grad() before the 2nd backward.
But there is no self.zero_grad() before self.criterion(self.model(samples), targe).backward() in line 77.

[rollovd]

Ciao @draym28!
Yeah, it makes sense.

I followed pseudo-code in LookSAM paper and I was thinking that we need to accumulate gradients under that condition. Therefore, I avoided step to zero gradients here.

Now, I rechecked SAM's code and they do zero gradients both times.
Apparently, you're right and I have a mistake here.

Have you ever tried this LookSAM solution with "under condition" zero_grad step, as you mentioned (perhaps you added it manually), or without it, as in my code? Does it work for you?

Thank you for being in touch haha!

[draym28]

I use your code in my project (thank you for the code, it is a big help!), and I find that with or without the zero_grad() step before the 2nd backward() will both bring benefits in different situations.

e.g., if you want to include the original loss as well ($L_{origin}+L_{perturbed}$) to compute the final gradient at every $k$ steps, you can avoid the zero_grad(), that will improve the performance in some cases, depending on your model design.

[rollovd]

Glad to hear it helps you! :)

Maybe in this case that's reasonable to add a boolean option "zero gradient before 2nd backward" :)
I'll do it right now.

Thank you for your contribution!
I hope project goes well hehe

[rollovd]

@draym28 How did you resolve an issue of "gv" key absence here while putting zero gradient after first backward?

Due to zero_grad after first backward, all the gradients will be None during next loop...

And we can't avoid condition with p.grad is None, because SAM has the same structure here, which makes sense.
I don't deny that I could have made a mistake in the code :)

[draym28]

@draym28 How did you resolve an issue of "gv" key absence here while putting zero gradient after first backward?

Due to zero_grad after first backward, all the gradients will be None during next loop... And as I understood it's okay, because SAM has the same line here. I don't deny that I could have made a mistake in the code :)

For LookSAM, the middle zero_grad() should only be conducted when if not k % self.k before 2n backward. So when k % self.k, the "gv" key still exist (not conduct zero_grad()), there will be no issue here.

[rollovd]

@draym28 Yeah, definitely. I made a typo :)

Thank you!

[rollovd]

@draym28 once again thank you for your contribution.
I updated code. Commit can be found here.

I guess this issue has resolved. Therefore, I can close this thread.