LSQ+的pytorch实现

import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function

class Round(Function):

    @staticmethod
    def forward(self, input):
        output = torch.round(input)
        return output

    @staticmethod
    def backward(self, grad_output):
        grad_input = grad_output.clone()
        return grad_input


def quant(x, scale,zeropoint):
    return Round.apply(torch.clamp((x-zeropoint) / scale, -128, 127))


def dequant(x, scale,zeropoint):
    return x * scale+zeropoint

def quantization_error(x,scale,zeropoint):
    xx=dequant(quant(x,scale,zeropoint),scale,zeropoint)
    return torch.sum((xx-x)**2)


class Quantization_error(nn.Module):
  def __init__(self):
    super().__init__()
    self.scale = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True)
    self.zeropoint = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True)
    self.scale=torch.nn.parameter.Parameter(torch.tensor(1.0))
    self.zeropoint =torch.nn.parameter.Parameter(torch.tensor(0.0))
  
  def forward(self,input):
    return quantization_error(input,self.scale,self.zeropoint)



def min_quantization_error(x):
    loss=0
    min_error=Quantization_error()
    min_error.to(device)
    optimizer = torch.optim.Adam(min_error.parameters(), lr=0.0001)
    for step in range(100000):
        x=x.to(device)
        pred=min_error(x)
        optimizer.zero_grad()
        pred.backward(retain_graph=True)
        optimizer.step()
        loss+=pred

        if step%1000==0 and step!=0:
            print(loss/1000)
            print(min_error.scale.item(),min_error.zeropoint.item())
            loss=0

    return min_error.scale,min_error.zeropoint

# ********************* 量化卷积（同时量化A/W，并做卷积） *********************
class Conv2d_Q(nn.Conv2d):
    def __init__(
            self,
            in_channels,
            out_channels,
            kernel_size,
            stride=1,
            padding=0,
            dilation=1,
            groups=1,
            bias=True,
            first_layer=0,
    ):
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias
        )
        self.weight_scale = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True)
        self.activation_scale = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True)
        self.activation_zeropoint = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True)
        self.fist_batch = 0
        self.first_layer = first_layer

    def forward(self, input):
        if self.fist_batch == 0:
            #初始化权重的scale
            Wmean=torch.mean(self.weight)
            Wstd=torch.std(self.weight)
            Wscale=torch.max(torch.abs(Wmean-3*Wstd),torch.abs(Wmean+3*Wstd))/128
            self.weight_scale = torch.nn.parameter.Parameter(Wscale)
            #初始化激活的scale和zeropoint
            Ascale,Azeropoint=min_quantization_error(input)
            self.activation_scale=torch.nn.parameter.Parameter(Ascale)
            self.activation_zeropoint=torch.nn.parameter.Parameter(Azeropoint)
            print("init finished")
            self.fist_batch = 1
        # 量化A和W
        if not self.first_layer:
            input = dequant(quant(input, self.activation_scale,self.activation_zeropoint), self.activation_scale,self.activation_zeropoint)
        q_input = input
        q_weight = dequant(quant(self.weight, self.weight_scale,0), self.weight_scale,0)
        # 量化卷积
        output = F.conv2d(
            input=q_input,
            weight=q_weight,
            bias=self.bias,
            stride=self.stride,
            padding=self.padding,
            dilation=self.dilation,
            groups=self.groups
        )
        return output


class QuanConv2d(nn.Module):
    def __init__(self, input_channels, output_channels,
                 kernel_size=-1, stride=-1, padding=-1, groups=1, last_relu=0, first_layer=0):
        super(QuanConv2d, self).__init__()
        self.last_relu = last_relu
        self.first_layer = first_layer
        self.q_conv = Conv2d_Q(input_channels, output_channels,
                               kernel_size=kernel_size, stride=stride, padding=padding, groups=groups,
                               first_layer=first_layer)
        self.bn = nn.BatchNorm2d(output_channels)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        if not self.first_layer:
            x = self.relu(x)
        x = self.q_conv(x)
        x = self.bn(x)
        if self.last_relu:
            x = self.relu(x)
        return x


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.quan_model = nn.Sequential(
            QuanConv2d(1, 8, kernel_size=3, stride=1, padding=1, first_layer=1),
            nn.MaxPool2d(kernel_size=2, stride=2),

            QuanConv2d(8, 16, kernel_size=3, stride=1, padding=1),
            QuanConv2d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),

            QuanConv2d(32, 10, kernel_size=3, stride=1, padding=1, last_relu=1),
            nn.AvgPool2d(kernel_size=7, stride=1, padding=0),
        )

    def forward(self, x):
        x = self.quan_model(x)
        x = x.view(x.size(0), -1)
        return x


import time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torchvision
import torchvision.transforms as transforms

device = torch.device('cuda:0')


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True

def train(epoch):
    model.train()

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        data, target = Variable(data), Variable(target)
        output = model(data)
        loss = criterion(output, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR: {}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.data.item(),
                optimizer.param_groups[0]['lr']))
    return


def test():
    model.eval()
    test_loss = 0
    correct = 0

    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        data, target = Variable(data), Variable(target)
        output = model(data)
        test_loss += criterion(output, target).data.item()
        pred = output.data.max(1, keepdim=True)[1]
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()
    acc = 100. * float(correct) / len(test_loader.dataset)

    print('acc is {}'.format(acc))


if __name__ == '__main__':
    setup_seed(int(time.time()))

    print('==> Preparing data..')
    train_dataset = torchvision.datasets.MNIST(root='../../data', train=True, transform=transforms.ToTensor(),
                                               download=True)
    test_dataset = torchvision.datasets.MNIST(root='../../data', train=False, transform=transforms.ToTensor())
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)

    print('******Initializing model******')
    model = Net()
    model.to(device)
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.xavier_uniform_(m.weight.data)
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            m.weight.data.normal_(0, 0.01)
            m.bias.data.zero_()


    criterion = nn.CrossEntropyLoss()


    base_lr = float(0.001)
    param_dict = dict(model.named_parameters())
    params = []
    for key, value in param_dict.items():
        if key=='quan_model.0.q_conv.weight_scale':
            g=1/torch.sqrt(torch.tensor(127.0*8*1*3*3))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.0.q_conv.activation_scale':
            g=1/torch.sqrt(torch.tensor(127.0*128*1*28*28))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.2.q_conv.weight_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*16*8*3*3))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.2.q_conv.activation_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*128*8*14*14))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.3.q_conv.weight_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*32*16*3*3))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.3.q_conv.activation_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*128*16*14*14))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.5.q_conv.weight_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*10*32*3*3))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        elif key=='quan_model.5.q_conv.activation_scale':
            g = 1 / torch.sqrt(torch.tensor(127.0*128*32*7*7))
            params += [{'params': [value], 'lr': base_lr*g, 'weight_decay': 0.00005}]
        else:
            params += [{'params': [value], 'lr': base_lr, 'weight_decay': 0.00005}]

    optimizer = optim.Adam(params, lr=base_lr, weight_decay=0.00005)

    for epoch in range(1, 30):
        train(epoch)
        test()
    

按照论文LSQ+: Improving low-bit quantization through learnable offsets and better initialization,权重采用的初始化方法为 而激活的初始化方法为 通过BP算法求得激活的 s i n i t 和 β i n i t s_{init}和\beta_{init} sinit​和βinit​