#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""ResNe(X)t models."""

from core.config import cfg
from models.blocks import (
    activation, 
    conv2d, 
    conv2d_cx, 
    gap2d, 
    gap2d_cx,
    init_weights,
    linear,
    linear_cx,
    norm2d, 
    norm2d_cx,
    pool2d,
    pool2d_cx,
)
from torch.nn import Module

#stage depths for ImageNet models
_IN_STAGE_DS ={50:(3,4,6,3),101:(3,4,23,3),152:(3,8,36,3)}


def get_trans_fun(name):
    """Retrieves the transformation function by name."""
    trans_funs={
        "basic_transform": BasicTransform,
        "bottleneck_transform":BottleneckTransform,
    }
    err_str="Transformation function '{}' not supported"
    assert name in trans_funs.keys(),err_str.format(name)
    return trans_funs[name]

class ResHead(Module):
    """ResNet head: AvgPool, 1x1."""

    def __init__(self, w_in, num_classes):
        super(ResHead, self).__init__()
        self.avg_pool = gap2d(w_in)
        self.fc = linear(w_in, num_classes, bias=True)

    def forward(self, x):
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
    
    @staticmethod
    def complexity(cx, w_in, num_classes):
        cx = gap2d_cx(cx, w_in)
        cx = linear_cx(cx, w_in, num_classes, bias=True)
        return cx


class BasicTransform(Module):
    """Basic transformation: 3x3,BN,AF,3x3,BN
     in artical :resnet-18/34 
    """

    def __init__(self,w_in,w_out,stride,w_b=None,groups=1):
        err_str = "Basic transform does not support w_b and groups options"
        assert w_b is None  and groups ==1,err_str
        super(BasicTransform, self).__init__()
        self.a=conv2d(w_in,w_out,3,stride=stride)
        self.a_bn=norm2d(w_out)
        self.a_af=activation()
        self.b=conv2d(w_out,w_out,3)
        self.b_bn=norm2d(w_out)
        self.b_bn.final_bn=True
    
    def forward(self,x):
        for layer in self.children():
            x=layer(x)
        return x

    @staticmethod
    def complexity(cx,w_in,w_out,stride,w_b=None,groups=1):
        err_str = "Basic transform does not support w_b and groups options"
        assert w_b is None and groups == 1, err_str
        cx = conv2d_cx(cx, w_in, w_out, 3, stride=stride)
        cx = norm2d_cx(cx, w_out)
        cx = conv2d_cx(cx, w_out, w_out, 3)
        cx = norm2d_cx(cx, w_out)
        return cx


class BottleneckTransform(Module):
    """
    Bottleneck transformation :1x1 ,BN,AF,3x3,BN,AF,1x1,BN.
    In artical used in resnet-50/101/159
    """
    def __init__(self,w_in,w_out,stride,w_b,groups=1):
        super(BottleneckTransform, self).__init__()
        # MSRA -> stride=2 is on 1x1; TH/C2 -> stride=2 is on 3x3
        (s1, s3) = (stride, 1) if cfg.RESNET.STRIDE_1X1 else (1, stride)
        self.a = conv2d(w_in, w_b, 1, stride=s1)
        self.a_bn = norm2d(w_b)
        self.a_af = activation()
        self.b = conv2d(w_b, w_b, 3, stride=s3, groups=groups)
        self.b_bn = norm2d(w_b)
        self.b_af = activation()
        self.c = conv2d(w_b, w_out, 1)
        self.c_bn = norm2d(w_out)
        self.c_bn.final_bn = True

    def forward(self, x):
        for layer in self.children():
            x = layer(x)
        return x

    @staticmethod
    def complexity(cx, w_in, w_out, stride, w_b, groups):
        (s1, s3) = (stride, 1) if cfg.RESNET.STRIDE_1X1 else (1, stride)
        cx = conv2d_cx(cx, w_in, w_b, 1, stride=s1)
        cx = norm2d_cx(cx, w_b)
        cx = conv2d_cx(cx, w_b, w_b, 3, stride=s3, groups=groups)
        cx = norm2d_cx(cx, w_b)
        cx = conv2d_cx(cx, w_b, w_out, 1)
        cx = norm2d_cx(cx, w_out)
        return cx

class ResBlock(Module):
    """Residual block: x+f(x)"""
    def __init__(self,w_in,w_out,stride,trans_fun,w_b=None,groups=1):
        super(ResBlock, self).__init__()
        self.proj,self.bn=None,None
        if(w_in!=w_out) or (stride!=1):
            self.proj=conv2d(w_in,w_out,1,stride=stride)
            self.bn=norm2d(w_out)
        self.f = trans_fun(w_in, w_out, stride, w_b, groups)
        self.af=activation()

    def forward(self,x):
        x_p=self.bn(self.proj(x)) if self.proj else x
        return self.af(x_p + self.f(x))
    
    @staticmethod
    def complexity(cx, w_in, w_out, stride, trans_fun, w_b, groups):
        if (w_in != w_out) or (stride != 1):
            h, w = cx["h"], cx["w"]
            cx = conv2d_cx(cx, w_in, w_out, 1, stride=stride)
            cx = norm2d_cx(cx, w_out)
            cx["h"], cx["w"] = h, w
        cx = trans_fun.complexity(cx, w_in, w_out, stride, w_b, groups)
        return cx

class ResStage(Module):
    """Stage of ResNet."""

    def __init__(self, w_in, w_out, stride, d, w_b=None, groups=1):
        super(ResStage, self).__init__()
        for i in range(d):
            b_stride=stride if i==0 else 1
            b_w_in=w_in if i==0 else w_out
            trans_fun=get_trans_fun(cfg.RESNET.TRANS_FUN)
            res_block=ResBlock(b_w_in,w_out,b_stride,trans_fun,w_b,groups)
            self.add_module("b{}".format(i+1),res_block)
        
    def forward(self,x):
        for block in self.children():
            x=block(x)
        return x

    @staticmethod
    def complexity(cx, w_in, w_out, stride, d, w_b=None, groups=1):
        for i in range(d):
            b_stride = stride if i == 0 else 1
            b_w_in = w_in if i == 0 else w_out
            trans_f = get_trans_fun(cfg.RESNET.TRANS_FUN)
            cx = ResBlock.complexity(cx, b_w_in, w_out, b_stride, trans_f, w_b, groups)
        return cx


class ResStemIN(Module):
    """ResNet stem for ImageNet: 7x7, BN, AF, MaxPool."""

    def __init__(self, w_in, w_out):
        super(ResStemIN, self).__init__()
        self.conv=conv2d(w_in,w_out,7,stride=2)
        self.bn=norm2d(w_out)
        self.af=activation()
        self.pool=pool2d(w_out,3,stride=2)
    
    def forward(self,x):
        for layer in self.children():
            x=layer(x)
        return x
    
    @staticmethod
    def complexity(cx, w_in, w_out):
        cx = conv2d_cx(cx, w_in, w_out, 7, stride=2)
        cx = norm2d_cx(cx, w_out)
        cx = pool2d_cx(cx, w_out, 3, stride=2)
        return cx
    

class ResNet(Module):
    """ResNet model."""
    def __init__(self):
        super(ResNet, self).__init__()
        self._construct_imagenet()
        self.apply(init_weights)

    def _construct_imagenet(self):
        g,gw=cfg.RESNET.NUM_GROUPS,cfg.RESNET.WIDTH_PER_GROUP
        (d1,d2,d3,d4)=_IN_STAGE_DS[cfg.MODEL.DEPTH]
        w_b=gw*g
        self.stem=ResStemIN(3,64)
        self.s1=ResStage(64,256,stride=1,d=d1,w_b=w_b,groups=g)
        self.s2=ResStage(256,512,stride=2,d=d2,w_b=w_b*2,groups=g)
        self.s3=ResStage(512,1024,stride=2,d=d3,w_b=w_b*4,groups=g)
        self.s4=ResStage(1024,2048,stride=2,d=d4,w_b=w_b*8,groups=g)
        self.head=ResHead(2048,cfg.MODEL.NUM_CLASSES)
    
    def forward(self,x):
        for module in self.children():
            x=module(x)
        return x
    
    @staticmethod
    def complexity(cx):
        """Computes model complexity. If you alter the model, make sure to update."""
        g, gw = cfg.RESNET.NUM_GROUPS, cfg.RESNET.WIDTH_PER_GROUP
        (d1, d2, d3, d4) = _IN_STAGE_DS[cfg.MODEL.DEPTH]
        w_b = gw * g
        cx = ResStemIN.complexity(cx, 3, 64)
        cx = ResStage.complexity(cx, 64, 256, 1, d=d1, w_b=w_b, groups=g)
        cx = ResStage.complexity(cx, 256, 512, 2, d=d2, w_b=w_b * 2, groups=g)
        cx = ResStage.complexity(cx, 512, 1024, 2, d=d3, w_b=w_b * 4, groups=g)
        cx = ResStage.complexity(cx, 1024, 2048, 2, d=d4, w_b=w_b * 8, groups=g)
        cx = ResHead.complexity(cx, 2048, cfg.MODEL.NUM_CLASSES)
        return cx

#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""ResNe(X)t models."""

from core.config import cfg
from models.blocks import (
    activation, 
    conv2d, 
    conv2d_cx, 
    gap2d, 
    gap2d_cx,
    init_weights,
    linear,
    linear_cx,
    norm2d, 
    norm2d_cx,
    pool2d,
    pool2d_cx,
)
from torch.nn import Module

#stage depths for ImageNet models
_IN_STAGE_DS ={50:(3,4,6,3),101:(3,4,23,3),152:(3,8,36,3)}


def get_trans_fun(name):
    """Retrieves the transformation function by name."""
    trans_funs={
        "basic_transform": BasicTransform,
        "bottleneck_transform":BottleneckTransform,
    }
    err_str="Transformation function '{}' not supported"
    assert name in trans_funs.keys(),err_str.format(name)
    return trans_funs[name]

class ResHead(Module):
    """ResNet head: AvgPool, 1x1."""

    def __init__(self, w_in, num_classes):
        super(ResHead, self).__init__()
        self.avg_pool = gap2d(w_in)
        self.fc = linear(w_in, num_classes, bias=True)

    def forward(self, x):
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
    
    @staticmethod
    def complexity(cx, w_in, num_classes):
        cx = gap2d_cx(cx, w_in)
        cx = linear_cx(cx, w_in, num_classes, bias=True)
        return cx


class BasicTransform(Module):
    """Basic transformation: 3x3,BN,AF,3x3,BN
     in artical :resnet-18/34 
    """

    def __init__(self,w_in,w_out,stride,w_b=None,groups=1):
        err_str = "Basic transform does not support w_b and groups options"
        assert w_b is None  and groups ==1,err_str
        super(BasicTransform, self).__init__()
        self.a=conv2d(w_in,w_out,3,stride=stride)
        self.a_bn=norm2d(w_out)
        self.a_af=activation()
        self.b=conv2d(w_out,w_out,3)
        self.b_bn=norm2d(w_out)
        self.b_bn.final_bn=True
    
    def forward(self,x):
        for layer in self.children():
            x=layer(x)
        return x

    @staticmethod
    def complexity(cx,w_in,w_out,stride,w_b=None,groups=1):
        err_str = "Basic transform does not support w_b and groups options"
        assert w_b is None and groups == 1, err_str
        cx = conv2d_cx(cx, w_in, w_out, 3, stride=stride)
        cx = norm2d_cx(cx, w_out)
        cx = conv2d_cx(cx, w_out, w_out, 3)
        cx = norm2d_cx(cx, w_out)
        return cx


class BottleneckTransform(Module):
    """
    Bottleneck transformation :1x1 ,BN,AF,3x3,BN,AF,1x1,BN.
    In artical used in resnet-50/101/159
    """
    def __init__(self,w_in,w_out,stride,w_b,groups=1):
        super(BottleneckTransform, self).__init__()
        # MSRA -> stride=2 is on 1x1; TH/C2 -> stride=2 is on 3x3
        (s1, s3) = (stride, 1) if cfg.RESNET.STRIDE_1X1 else (1, stride)
        self.a = conv2d(w_in, w_b, 1, stride=s1)
        self.a_bn = norm2d(w_b)
        self.a_af = activation()
        self.b = conv2d(w_b, w_b, 3, stride=s3, groups=groups)
        self.b_bn = norm2d(w_b)
        self.b_af = activation()
        self.c = conv2d(w_b, w_out, 1)
        self.c_bn = norm2d(w_out)
        self.c_bn.final_bn = True

    def forward(self, x):
        for layer in self.children():
            x = layer(x)
        return x

    @staticmethod
    def complexity(cx, w_in, w_out, stride, w_b, groups):
        (s1, s3) = (stride, 1) if cfg.RESNET.STRIDE_1X1 else (1, stride)
        cx = conv2d_cx(cx, w_in, w_b, 1, stride=s1)
        cx = norm2d_cx(cx, w_b)
        cx = conv2d_cx(cx, w_b, w_b, 3, stride=s3, groups=groups)
        cx = norm2d_cx(cx, w_b)
        cx = conv2d_cx(cx, w_b, w_out, 1)
        cx = norm2d_cx(cx, w_out)
        return cx

class ResBlock(Module):
    """Residual block: x+f(x)"""
    def __init__(self,w_in,w_out,stride,trans_fun,w_b=None,groups=1):
        super(ResBlock, self).__init__()
        self.proj,self.bn=None,None
        if(w_in!=w_out) or (stride!=1):
            self.proj=conv2d(w_in,w_out,1,stride=stride)
            self.bn=norm2d(w_out)
        self.f = trans_fun(w_in, w_out, stride, w_b, groups)
        self.af=activation()

    def forward(self,x):
        x_p=self.bn(self.proj(x)) if self.proj else x
        return self.af(x_p + self.f(x))
    
    @staticmethod
    def complexity(cx, w_in, w_out, stride, trans_fun, w_b, groups):
        if (w_in != w_out) or (stride != 1):
            h, w = cx["h"], cx["w"]
            cx = conv2d_cx(cx, w_in, w_out, 1, stride=stride)
            cx = norm2d_cx(cx, w_out)
            cx["h"], cx["w"] = h, w
        cx = trans_fun.complexity(cx, w_in, w_out, stride, w_b, groups)
        return cx

class ResStage(Module):
    """Stage of ResNet."""

    def __init__(self, w_in, w_out, stride, d, w_b=None, groups=1):
        super(ResStage, self).__init__()
        for i in range(d):
            b_stride=stride if i==0 else 1
            b_w_in=w_in if i==0 else w_out
            trans_fun=get_trans_fun(cfg.RESNET.TRANS_FUN)
            res_block=ResBlock(b_w_in,w_out,b_stride,trans_fun,w_b,groups)
            self.add_module("b{}".format(i+1),res_block)
        
    def forward(self,x):
        for block in self.children():
            x=block(x)
        return x

    @staticmethod
    def complexity(cx, w_in, w_out, stride, d, w_b=None, groups=1):
        for i in range(d):
            b_stride = stride if i == 0 else 1
            b_w_in = w_in if i == 0 else w_out
            trans_f = get_trans_fun(cfg.RESNET.TRANS_FUN)
            cx = ResBlock.complexity(cx, b_w_in, w_out, b_stride, trans_f, w_b, groups)
        return cx


class ResStemIN(Module):
    """ResNet stem for ImageNet: 7x7, BN, AF, MaxPool."""

    def __init__(self, w_in, w_out):
        super(ResStemIN, self).__init__()
        self.conv=conv2d(w_in,w_out,7,stride=2)
        self.bn=norm2d(w_out)
        self.af=activation()
        self.pool=pool2d(w_out,3,stride=2)
    
    def forward(self,x):
        for layer in self.children():
            x=layer(x)
        return x
    
    @staticmethod
    def complexity(cx, w_in, w_out):
        cx = conv2d_cx(cx, w_in, w_out, 7, stride=2)
        cx = norm2d_cx(cx, w_out)
        cx = pool2d_cx(cx, w_out, 3, stride=2)
        return cx
    

class ResNet(Module):
    """ResNet model."""
    def __init__(self):
        super(ResNet, self).__init__()
        self._construct_imagenet()
        self.apply(init_weights)

    def _construct_imagenet(self):
        g,gw=cfg.RESNET.NUM_GROUPS,cfg.RESNET.WIDTH_PER_GROUP
        (d1,d2,d3,d4)=_IN_STAGE_DS[cfg.MODEL.DEPTH]
        w_b=gw*g
        self.stem=ResStemIN(3,64)
        self.s1=ResStage(64,256,stride=1,d=d1,w_b=w_b,groups=g)
        self.s2=ResStage(256,512,stride=2,d=d2,w_b=w_b*2,groups=g)
        self.s3=ResStage(512,1024,stride=2,d=d3,w_b=w_b*4,groups=g)
        self.s4=ResStage(1024,2048,stride=2,d=d4,w_b=w_b*8,groups=g)
        self.head=ResHead(2048,cfg.MODEL.NUM_CLASSES)
    
    def forward(self,x):
        for module in self.children():
            x=module(x)
        return x
    
    @staticmethod
    def complexity(cx):
        """Computes model complexity. If you alter the model, make sure to update."""
        g, gw = cfg.RESNET.NUM_GROUPS, cfg.RESNET.WIDTH_PER_GROUP
        (d1, d2, d3, d4) = _IN_STAGE_DS[cfg.MODEL.DEPTH]
        w_b = gw * g
        cx = ResStemIN.complexity(cx, 3, 64)
        cx = ResStage.complexity(cx, 64, 256, 1, d=d1, w_b=w_b, groups=g)
        cx = ResStage.complexity(cx, 256, 512, 2, d=d2, w_b=w_b * 2, groups=g)
        cx = ResStage.complexity(cx, 512, 1024, 2, d=d3, w_b=w_b * 4, groups=g)
        cx = ResStage.complexity(cx, 1024, 2048, 2, d=d4, w_b=w_b * 8, groups=g)
        cx = ResHead.complexity(cx, 2048, cfg.MODEL.NUM_CLASSES)
        return cx

#!/usr/bin/env python3 # Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. ResNe(X)t models. #stage depths for ImageNet models Retrieves the transformation function by name. ResNet head: AvgPool, 1x1. Basic transformation: 3x3,BN,AF,3x3,BN      in artical :resnet-18/34 Bottleneck transformation :1x1 ,BN,AF,3x3,BN,AF,1x1,BN.     In artical used in resnet-50/101/159 # MSRA -> stride=2 is on 1x1; TH/C2 -> stride=2 is on 3x3 Residual block: x+f(x) Stage of ResNet. ResNet stem for ImageNet: 7x7, BN, AF, MaxPool. ResNet model. Computes model complexity. If you alter the model, make sure to update.