q-future/Co-Instruct

1# ------------------------------------------------------------------------2# Copyright (c) 2022 megvii-model. All Rights Reserved.3# ------------------------------------------------------------------------4# Source: https://github.com/megvii-research/NAFNet5 6import numpy as np7import torch8import torch.nn as nn9import torch.nn.functional as F10import math11 12class LayerNormFunction(torch.autograd.Function):13 14    @staticmethod15    def forward(ctx, x, weight, bias, eps):16        ctx.eps = eps17        N, C, H, W = x.size()18        mu = x.mean(1, keepdim=True)19        var = (x - mu).pow(2).mean(1, keepdim=True)20        y = (x - mu) / (var + eps).sqrt()21        ctx.save_for_backward(y, var, weight)22        y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)23        return y24 25    @staticmethod26    def backward(ctx, grad_output):27        eps = ctx.eps28 29        N, C, H, W = grad_output.size()30        y, var, weight = ctx.saved_variables31        g = grad_output * weight.view(1, C, 1, 1)32        mean_g = g.mean(dim=1, keepdim=True)33 34        mean_gy = (g * y).mean(dim=1, keepdim=True)35        gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)36        return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum(37            dim=0), None38 39class LayerNorm2d(nn.Module):40 41    def __init__(self, channels, eps=1e-6):42        super(LayerNorm2d, self).__init__()43        self.register_parameter('weight', nn.Parameter(torch.ones(channels)))44        self.register_parameter('bias', nn.Parameter(torch.zeros(channels)))45        self.eps = eps46 47    def forward(self, x):48        return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)49    50 51 52class AvgPool2d(nn.Module):53    def __init__(self, kernel_size=None, base_size=None, auto_pad=True, fast_imp=False, train_size=None):54        super().__init__()55        self.kernel_size = kernel_size56        self.base_size = base_size57        self.auto_pad = auto_pad58 59        # only used for fast implementation60        self.fast_imp = fast_imp61        self.rs = [5, 4, 3, 2, 1]62        self.max_r1 = self.rs[0]63        self.max_r2 = self.rs[0]64        self.train_size = train_size65 66    def extra_repr(self) -> str:67        return 'kernel_size={}, base_size={}, stride={}, fast_imp={}'.format(68            self.kernel_size, self.base_size, self.kernel_size, self.fast_imp69        )70 71    def forward(self, x):72        if self.kernel_size is None and self.base_size:73            train_size = self.train_size74            if isinstance(self.base_size, int):75                self.base_size = (self.base_size, self.base_size)76            self.kernel_size = list(self.base_size)77            self.kernel_size[0] = x.shape[2] * self.base_size[0] // train_size[-2]78            self.kernel_size[1] = x.shape[3] * self.base_size[1] // train_size[-1]79 80            # only used for fast implementation81            self.max_r1 = max(1, self.rs[0] * x.shape[2] // train_size[-2])82            self.max_r2 = max(1, self.rs[0] * x.shape[3] // train_size[-1])83 84        if self.kernel_size[0] >= x.size(-2) and self.kernel_size[1] >= x.size(-1):85            return F.adaptive_avg_pool2d(x, 1)86 87        if self.fast_imp:  # Non-equivalent implementation but faster88            h, w = x.shape[2:]89            if self.kernel_size[0] >= h and self.kernel_size[1] >= w:90                out = F.adaptive_avg_pool2d(x, 1)91            else:92                r1 = [r for r in self.rs if h % r == 0][0]93                r2 = [r for r in self.rs if w % r == 0][0]94                # reduction_constraint95                r1 = min(self.max_r1, r1)96                r2 = min(self.max_r2, r2)97                s = x[:, :, ::r1, ::r2].cumsum(dim=-1).cumsum(dim=-2)98                n, c, h, w = s.shape99                k1, k2 = min(h - 1, self.kernel_size[0] // r1), min(w - 1, self.kernel_size[1] // r2)100                out = (s[:, :, :-k1, :-k2] - s[:, :, :-k1, k2:] - s[:, :, k1:, :-k2] + s[:, :, k1:, k2:]) / (k1 * k2)101                out = torch.nn.functional.interpolate(out, scale_factor=(r1, r2))102        else:103            n, c, h, w = x.shape104            s = x.cumsum(dim=-1).cumsum_(dim=-2)105            s = torch.nn.functional.pad(s, (1, 0, 1, 0))  # pad 0 for convenience106            k1, k2 = min(h, self.kernel_size[0]), min(w, self.kernel_size[1])107            s1, s2, s3, s4 = s[:, :, :-k1, :-k2], s[:, :, :-k1, k2:], s[:, :, k1:, :-k2], s[:, :, k1:, k2:]108            out = s4 + s1 - s2 - s3109            out = out / (k1 * k2)110 111        if self.auto_pad:112            n, c, h, w = x.shape113            _h, _w = out.shape[2:]114            # print(x.shape, self.kernel_size)115            pad2d = ((w - _w) // 2, (w - _w + 1) // 2, (h - _h) // 2, (h - _h + 1) // 2)116            out = torch.nn.functional.pad(out, pad2d, mode='replicate')117 118        return out119 120def replace_layers(model, base_size, train_size, fast_imp, **kwargs):121    for n, m in model.named_children():122        if len(list(m.children())) > 0:123            ## compound module, go inside it124            replace_layers(m, base_size, train_size, fast_imp, **kwargs)125 126        if isinstance(m, nn.AdaptiveAvgPool2d):127            pool = AvgPool2d(base_size=base_size, fast_imp=fast_imp, train_size=train_size)128            assert m.output_size == 1129            setattr(model, n, pool)130 131 132'''133ref. 134@article{chu2021tlsc,135  title={Revisiting Global Statistics Aggregation for Improving Image Restoration},136  author={Chu, Xiaojie and Chen, Liangyu and and Chen, Chengpeng and Lu, Xin},137  journal={arXiv preprint arXiv:2112.04491},138  year={2021}139}140'''141class Local_Base():142    def convert(self, *args, train_size, **kwargs):143        replace_layers(self, *args, train_size=train_size, **kwargs)144        imgs = torch.rand(train_size)145        with torch.no_grad():146            self.forward(imgs)