def forward()

in Classification/models/deconv.py [0:0]

• 68 lines of code
• 18 McCabe index (conditional complexity)

    def forward(self, x):

        if x.numel()==0:
            return x
        N, C, H, W = x.shape
        B = self.block

        x=x.contiguous()
        if self.norm_type=='l1norm':
            x_norm=x.abs().mean(dim=(1,2,3),keepdim=True)            
            x =  x/ (x_norm + self.eps)

        elif self.norm_type=='layernorm':
            x=self.layernorm(x)


        if self.training:
            self.counter+=1
            frozen=self.freeze and (self.counter% (self.freeze_iter * 10) >self.freeze_iter)

        if self.training and (not frozen):

            # 1. im2col: N x cols x pixels -> N*pixles x cols
            if self.kernel_size[0]>1:
                X = torch.nn.functional.unfold(x, self.kernel_size,self.dilation,self.padding,self.sampling_stride).transpose(1, 2).contiguous()                    
            else:
                #channel wise
                X = x.permute(0, 2, 3, 1).contiguous().view(-1, C)[::self.sampling_stride**2,:]

            if self.groups==1:
                # (C//B*N*pixels,k*k*B)
                X = X.view(-1, self.num_features, C // B).transpose(1, 2).contiguous().view(-1, self.num_features)
            else:
                X=X.view(-1,X.shape[-1])

            # 2. calculate mean,cov,cov_isqrt
            X_mean = X.mean(0)

            if self.groups==1:
                M = X.shape[0]
                XX_mean = X.t()@X/M 
            else:
                M = X.shape[1]
                XX_mean = X.transpose(1,2)@X/M 
                
            if self.sync:
                process_group = self.process_group
                world_size = 1
                if not self.process_group:
                    process_group = torch.distributed.group.WORLD
        
                if torch.distributed.is_initialized():
                    world_size = torch.distributed.get_world_size(process_group)
                    #sync once implementation:
                    sync_data=torch.cat([X_mean.view(-1),XX_mean.view(-1)],dim=0)
                    sync_data_list=[torch.empty_like(sync_data) for k in range(world_size)]
                    sync_data_list = diffdist.functional.all_gather(sync_data_list, sync_data)
                    sync_data=torch.stack(sync_data_list).mean(0)
                    X_mean=sync_data[:X_mean.numel()].view(X_mean.shape)
                    XX_mean=sync_data[X_mean.numel():].view(XX_mean.shape)
                

            if self.groups==1:
                Id = torch.eye(XX_mean.shape[1], dtype=X.dtype, device=X.device)
                cov= XX_mean- X_mean.unsqueeze(1) @X_mean.unsqueeze(0)+self.eps*Id
                cov_isqrt = isqrt_newton_schulz_autograd(cov, self.n_iter)
            else:
                Id = torch.eye(self.num_features, dtype=X.dtype, device=X.device).expand(self.groups, self.num_features, self.num_features)
                cov= (XX_mean- (X_mean.unsqueeze(2)) @X_mean.unsqueeze(1))+self.eps*Id
                cov_isqrt = isqrt_newton_schulz_autograd_batch(cov, self.n_iter)

            # track stats for evaluation.
            self.running_mean.mul_(1 - self.momentum)                
            self.running_mean.add_(X_mean.detach() * self.momentum)
            self.running_cov_isqrt.mul_(1 - self.momentum)
            self.running_cov_isqrt.add_(cov_isqrt.detach() * self.momentum)

        else:
            X_mean = self.running_mean
            cov_isqrt = self.running_cov_isqrt

        #3. S * X * D * W = (S * X) * (D * W)
        if self.groups==1:
            w = self.weight.view(-1, self.num_features, C // B).transpose(1, 2).contiguous().view(-1,self.num_features) @ cov_isqrt
            b = self.bias - (w @ (X_mean.unsqueeze(1))).view(self.weight.shape[0], -1).sum(1)
            w = w.view(-1, C // B, self.num_features).transpose(1, 2).contiguous()
        else:
            w = self.weight.view(C//B, -1,self.num_features)@cov_isqrt
            b = self.bias - (w @ (X_mean.view( -1,self.num_features,1))).view(self.bias.shape)

        w = w.view(self.weight.shape)
        
        x= F.conv2d(x, w, b, self.stride, self.padding, self.dilation, self.groups)

        return x