完整部署快照

ixformer_sdk/train/functions/group_norm.py

@@ -0,0 +1,61 @@
+import ixformer._C as ops
+import torch
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+
+class GN_NHWC_Func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, X: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, G: int, eps: float, activation: str):
+        X_out, means, rstds = ops.train.gn_nhwc_fwd(X, weight, bias, G, eps, activation)
+        ctx.save_for_backward(X, weight, bias, means, rstds)
+        ctx.G = G
+        ctx.activation = activation
+        return X_out
+
+    @staticmethod
+    def backward(ctx, dy: torch.Tensor):
+        dy = dy.contiguous(memory_format=torch.channels_last)
+        X, weight, bias, means, rstds = ctx.saved_tensors 
+        dx, dgamma, dbeta = ops.train.gn_nhwc_bwd(dy, X, weight, bias, means, rstds, ctx.G, ctx.activation)
+        return dx, dgamma, dbeta, None, None, None
+    
+
+class GroupNorm_nhwc(torch.nn.GroupNorm):
+    def __init__(self, num_groups: int, nc: int, activation='identity', **kwargs):
+        super().__init__(num_groups, nc, **kwargs)
+        assert activation in {'identity', 'silu', 'relu', 'gelu', 'gelu_tanh'}
+        if activation == 'identity':
+            self.activation = 0
+        if activation == 'relu':
+            self.activation = 1
+        if activation == 'silu':
+            self.activation = 2
+        if activation == 'gelu':
+            self.activation = 3
+        if activation == 'gelu_tanh':
+            self.activation = 4
+
+    @torch._dynamo.disable
+    def forward(self, x):
+        #print(x.shape, self.num_channels)
+        if len(x.size()) == 3:
+            N, C, L = x.shape
+        elif len(x.size()) == 4:
+            N, C, H, W = x.shape
+        else:
+            raise ValueError
+        G = self.num_groups
+
+        #if C // G > 512:
+        #    raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: C // G = {C // G} which is greater than 512. This input is not supported.')
+
+        #if H * W % 8 != 0:
+        #    raise ValueError(f'Error in fwd for X.shape={x.shape}, G={G}: H * W is not a multiple of 8. This input is not supported.')
+
+        if self.affine:
+            return GN_NHWC_Func.apply(x, self.weight, self.bias, self.num_groups, self.eps, self.activation)
+        else:
+            w = torch.ones((self.num_channels,), device=x.device, dtype=x.dtype)
+            b = torch.zeros((self.num_channels,), device=x.device, dtype=x.dtype)
+            return GN_NHWC_Func.apply(x, w, b, self.num_groups, self.eps, self.activation)