dydjw9 / efficient_sam Goto Github PK

BTW, there is a minor typo in ESAM init.

Could you please provide a demo code for training VIT?

Based on your implementation Esam.py, same parameters are randomly masked at the descent step (i.e., the second step as follows).
But in your paper (the paragraph above eq(5)), the mask is added to the perturbance $\epsilon$.
As they are very different, could you please provide some explanations?

@torch.no_grad()
    def second_step(self, zero_grad=False):
        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None or not self.state[p]: continue
                p.sub_(self.state[p]["e_w"])  # get back to "w" from "w + e(w)"
                self.state[p]["e_w"] = 0

                if random.random() > self.beta:
                    p.requires_grad = False

        self.base_optimizer.step()  # do the actual "sharpness-aware" update

        if zero_grad: self.zero_grad()

The line 26 of Esam.py is different from the original algorithm 1 proposed in the paper, i.e., beta in the code vs. (1 - beta) in the paper.

I came across your work and you did a fantastic job in improving the performance of SAM.

It seems that the current implementation supports only a single loss function.
While the code example does include the case for fp16, there are no mentions of gradient scaling, which is commonly used together in AMP.

Are there any plans to support multiple loss functions and GradScaler?