_beta

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function, absolute_import

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
import math
import torch.utils.model_zoo as model_zoo
import pdb
import numpy as np
from collections import OrderedDict



__all__ = ['ResNet', 'resnet20_ste_6']


model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
}

import torch
import torch.nn as nn
import torch.nn.functional as F

bit_num = 1
range_bit = 2 ** bit_num

class RoundWithGradient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x.round()
    @staticmethod
    def backward(ctx, g):
        return g

class GradientScale(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return x
    @staticmethod
    def backward(ctx, g):
        return g / 0.001

class ste(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        x_abs = torch.abs(x) + 0.5
        x_abs = torch.clamp(x_abs, min=0.5+(1e-4), max=(range_bit/2)+0.5-(1e-4))
        x_sign = torch.sign(x)
        x_sign = torch.sign(x_sign - (1e-4))

        output = x_sign * torch.round(x_abs)
        return output
    @staticmethod
    def backward(ctx, g):
        return g

class DSQConv(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True,
                momentum = 0.1,
                num_bit = bit_num, QInput = True, bSetQ = True):
        super(DSQConv, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.num_bit = num_bit
        self.quan_input = QInput
        self.bit_range = 2**self.num_bit - 1
        self.is_quan = bSetQ
        self.temp = -1
        self.sig_w = 5 # T:10 - 5, T:1 - 1
        self.sig_a = 5 # T:10 - 5, T:1 - 1
        self.q_value = torch.from_numpy(np.linspace(0,self.bit_range,self.bit_range+1))
        self.q_value = self.q_value.reshape(len(self.q_value),1,1,1,1).float().cuda()

        if self.is_quan:
            # using int32 max/min as init and backprogation to optimization
            # Weight
            self.uW = nn.Parameter(data = torch.tensor(2 **31 - 1).float().cuda())
            self.lW  = nn.Parameter(data = torch.tensor((-1) * (2**32)).float().cuda())
            self.register_buffer('init', torch.tensor(1).float().cuda())
            self.beta_w = nn.Parameter(data = torch.tensor(0.2).float().cuda())
            # Bias
            if self.bias is not None:
                self.uB = nn.Parameter(data = torch.tensor(2 **31 - 1).float())
                self.lB  = nn.Parameter(data = torch.tensor((-1) * (2**32)).float())
                self.register_buffer('running_uB', torch.tensor([self.uB.data]))# init with ub
                self.register_buffer('running_lB', torch.tensor([self.lB.data]))# init with lb
                self.alphaB = nn.Parameter(data = torch.tensor(1).float())

            # Activation input
            if self.quan_input:
                self.uA = nn.Parameter(data = torch.tensor(2 **31 - 1).float().cuda())
                self.lA  = nn.Parameter(data = torch.tensor((-1) * (2**32)).float().cuda())
                self.beta_a = nn.Parameter(data = torch.tensor(0.2).float().cuda())

    def clipping(self, x, upper, lower):
        # clip lower
        x = x + F.relu(lower - x)
        # clip upper
        x = x - F.relu(x - upper)

        return x

    def step(self, x):
        if self.num_bit == 1:
            output = (x+1) / 2
            output = RoundWithGradient.apply(output)
            return 2 * output - 1
        else:
            return RoundWithGradient.apply(x)

    def w_quan(self, x, u, l):
        delta = (u - l) / (self.bit_range)
        interval = (x - l) / delta
        output = 2 * RoundWithGradient.apply(interval) - self.bit_range
        return output

    def a_quan(self, x, u, l):
        delta = (u - l) / (self.bit_range)
        interval = (x - l) / delta
        output = RoundWithGradient.apply(interval)
        return output

    def sigmoid(self, x, T=2):
        output = 1 / (1+torch.exp(-(x)*T))
        return output

    def forward(self, x):

        if self.is_quan:
            if self.init:
                print(self.init)
                self.init = torch.tensor(0)
                self.lW.data = -self.weight.std() * 3
                self.uW.data = self.weight.std() * 3
                self.lA.data = -x.std() * 3
                self.uA.data = x.std() * 3
                self.beta_w.data = torch.mean(torch.abs(self.weight)) / self.bit_range
                self.beta_a.data = torch.mean(torch.abs(x)) / self.bit_range

            self.lW.data = torch.clamp(self.lW.data, min=-1e+1, max=-1e-6)
            self.uW.data = torch.clamp(self.uW.data, min=1e-6, max=1e+1)
            self.lA.data = torch.clamp(self.lA.data, min=-1e+1, max=-1e-6)
            self.uA.data = torch.clamp(self.uA.data, min=1e-6, max=1e+1)

            curr_running_lw = self.lW
            curr_running_uw = self.uW

            curr_running_la = 0
            curr_running_ua = self.uA

            # Weight kernel_soft_argmax
            Qweight = self.clipping(self.weight, curr_running_uw, curr_running_lw)
            Qweight = self.w_quan(Qweight, curr_running_uw, curr_running_lw)
            Qweight = torch.abs(self.beta_w) * Qweight
            Qbias = self.bias

            # Input(Activation)
            Qactivation = x
            if self.quan_input:
                Qactivation = self.clipping(x, curr_running_ua, curr_running_la)
                Qactivation = self.a_quan(Qactivation, curr_running_ua, curr_running_la)
                Qactivation = torch.abs(self.beta_a) * Qactivation

            output = F.conv2d(Qactivation, Qweight, Qbias,  self.stride, self.padding, self.dilation, self.groups)

        else:
            output =  F.conv2d(x, self.weight, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)

        return output

class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = DSQConv(inplanes, planes, kernel_size=3, stride=stride,
                             padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = DSQConv(planes, planes, kernel_size=3, stride=1,
                             padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.shortcut = nn.Sequential()
        if stride != 1 or inplanes != planes:
            self.shortcut = LambdaLayer(lambda x:
                                        F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes//4, planes//4), "constant", 0))

    def forward(self, x):
        conv1_out = F.relu(self.bn1(self.conv1(x)))
        conv2_out = self.bn2(self.conv2(conv1_out))
        out = conv2_out +  self.shortcut(x)
        out = F.relu(out)
        return out, conv1_out, conv2_out
#        conv1_out = self.conv1(self.bn1(x))
#        conv2_out = self.conv2(self.bn2(F.relu(conv1_out)))
#        out = conv2_out + self.shortcut(x)
#        out = F.relu(out)
#        return out, conv1_out, conv2_out

class ResNet(nn.Module):

    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 16
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)

        self.bn2 = nn.BatchNorm1d(64)
        # self.bn3 = nn.BatchNorm1d(num_classes)

        self.linear = nn.Linear(64, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        ret_dict = dict()

        for i, stride in enumerate(strides):
            layers = []
            layers.append(block(self.in_planes, planes, stride))
            ret_dict['block_{}'.format(i)] = nn.Sequential(*layers)
            self.in_planes = planes * block.expansion

        return nn.Sequential(OrderedDict(ret_dict))

    def forward(self, x):
        ret_dict = dict()
        # out = F.relu(self.bn1(self.conv1(x)))
        # out = F.hardtanh(self.bn1(self.conv1(x)))
        out = F.relu(self.conv1(x))
        layer_names = self.layer1._modules.keys()
        for i, layer_name in enumerate(layer_names):
            out, conv1_out, conv2_out = self.layer1._modules[layer_name](out)
            ret_dict['layer1_{}_conv1'.format(i)] = conv1_out
            ret_dict['layer1_{}_conv2'.format(i)] = conv2_out

        layer_names = self.layer2._modules.keys()
        for i, layer_name in enumerate(layer_names):
            out, conv1_out, conv2_out = self.layer2._modules[layer_name](out)
            ret_dict['layer2_{}_conv1'.format(i)] = conv1_out
            ret_dict['layer2_{}_conv2'.format(i)] = conv2_out

        layer_names = self.layer3._modules.keys()
        for i, layer_name in enumerate(layer_names):
            out, conv1_out, conv2_out = self.layer3._modules[layer_name](out)
            ret_dict['layer3_{}_conv1'.format(i)] = conv1_out
            ret_dict['layer3_{}_conv2'.format(i)] = conv2_out

        out = F.avg_pool2d(out, out.size()[3])
        out = out.view(out.size(0), -1)
        out = self.bn2(out)
        out = self.linear(out)
        ret_dict['out'] = out
        return ret_dict

def resnet20_ste_6(pretrained=True, **kwargs):
    """Constructs a ResNet-18 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    return ResNet(BasicBlock, [3, 3, 3], **kwargs)