Merge pull request #225 from kai422/main

correct ish_reshaper

openood/trainers/ish_trainer.py

@@ -98,26 +98,32 @@ def forward(
         x: torch.Tensor,
         weight: nn.Parameter,
         bias: nn.Parameter,
-        dropiter
+        ish_reshaper
     ):
-        ctx.dropiter = dropiter
+        ctx.ish_reshaper = ish_reshaper
         ctx.x_shape = x.shape
         ctx.has_bias = bias is not None
-        ctx.save_for_backward(dropiter.select(x, ctx), weight)
+        ctx.save_for_backward(ish_reshaper.select(x, ctx), weight)
         return F.linear(x, weight, bias)
 
     @staticmethod 
     def backward(ctx, grad_output: torch.Tensor):
         x, weight = ctx.saved_tensors
         grad_bias = torch.sum(grad_output, list(range(grad_output.dim()-1))) if ctx.has_bias else None
         ic, oc = weight.shape
-        x = ctx.dropiter.pad(x, ctx)
+        x = ctx.ish_reshaper.pad(x, ctx)
         grad_weight = grad_output.view(-1,ic).T.mm(x.view(-1,oc))
         grad_input = torch.matmul(grad_output, weight, out=x.view(ctx.x_shape))
         return grad_input, grad_weight, grad_bias, None
 
-linear_forward = _ISHTLinear.apply
+_linear_forward = _ISHTLinear.apply
 
+def linear_forward(self, x):
+    if self.training:
+        x = _linear_forward(x, self.weight, self.bias, self.ish_reshaper)
+    else:
+        x = F.linear(x, self.weight, self.bias)
+    return x
 
 supports = {
     nn.Linear: linear_forward,
@@ -152,7 +158,6 @@ def cache_minksample_expscale(self, x: torch.Tensor, ctx=None):
         return x
 
     def load_minksample_expscale(self, x, ctx=None):      
-        print(ctx.idxs.shape, x.shape, )
         return torch.zeros(
             ctx.shape, device=x.device, dtype=x.dtype
         ).scatter_(1, ctx.idxs, x)
@@ -217,7 +222,6 @@ def load_minksample_lnscale(self, x, ctx=None):
     def transfer(model, strategy, gamma, autocast):
         _type = type(model)
         ish_reshaper = ISHReshaper(strategy, gamma)
-        ish_reshaper.autocast = autocast # just for recording
         model.forward = partial(supports[_type], model)
         model.ish_reshaper = ish_reshaper
         print(f"{_type}.forward => ish.{strategy}.{_type}.forward")