<jupyter_start><jupyter_text>CNN Quantization PracticeThis tutorial implements [DoReFa-Net](https://arxiv.org/abs/1606.06160) weight and activation quantization scheme. Gradient quantization is not implemented, however is described in the original paper.Code inspiration is taken from randomly found [github](https://github.com/zzzxxxttt/pytorch_DoReFaNet) due to the fact that current notebook author has a lack of creativity.<jupyter_code>import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import numpy as np
from tqdm.auto import tqdm
from matplotlib import pyplot as plt
use_cuda = True
torch.manual_seed(42)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
batch_size = 256
test_batch_size = 128
epochs = 10
lr = 1e-3
weight_decay = 0<jupyter_output><empty_output><jupyter_text>Let's use MNIST dataset and evaluate one of the current state-of-the-art methods of CNN quantization.We will implement a simple CNN model yet capable of reaching high performance on MNIST.Then we will implement a quantized CNN capable of reaching same performance on it.  Data loadingWe are developing quantized neural networks, so we need to feed images into CNN in different manner:pixel absolute values are no longer needed to be less than 1.<jupyter_code>def to_signed_char(tensor):
    # note that we convert image values from float32 range([-0.5, 0.5]) to signed char range [-128,127]
    return torch.floor(255.0 * tensor)

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.5,), (1, )),
                       
                   ])),
    batch_size=batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (1, )),
                   ])),
    batch_size=test_batch_size, shuffle=True)<jupyter_output><empty_output><jupyter_text>Training and Validation Functions<jupyter_code>def train(model, device, train_loader, optimizer, epoch, **kwargs):
    log_interval = kwargs.get('log_interval', 10)
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(to_signed_char(data))
        loss = F.nll_loss(output, target)
        loss.backward(retain_graph=True)
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))


def test(model, device, test_loader, **kwargs):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(to_signed_char(data))
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))<jupyter_output><empty_output><jupyter_text>CNN ModelLet's implement something really easy, __conv->bn->relu->pool__ will be enough.<jupyter_code>class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1),
            nn.BatchNorm2d(32),
            nn.ReLU()
        )
        self.maxpool_2x2 = nn.MaxPool2d((2,2))
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        self.fc1 = nn.Linear(64, 10)
        self.gmp = nn.AdaptiveMaxPool2d(1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool_2x2(x)
        x = self.conv2(x)
        x = self.gmp(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        output = F.log_softmax(x, dim=1)
        return output
model = Net().to(device)
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
scheduler = StepLR(optimizer, step_size=1, gamma=0.9)
for epoch in range(1, epochs + 1):
    train(model, device, train_loader, optimizer, epoch, log_interval=50)
    test(model, device, test_loader)
    scheduler.step()<jupyter_output><empty_output><jupyter_text>Quantization LayersNot let's implement basic building blocks for DoReFa-Net quantization scheme.Firstly, we will need to implement uniform quantizer.Given real value $r\in [0;1]$ uniform quantized value $r_0$ may be calculated:$r_0 = \frac{1}{2^k-1} round((2^k-1) * r)$The calculated value may be saved as floating point number, or converted to fixed point representation by multiplying it on $2^{\text{n_bits}}$<jupyter_code>def get_uniform_quantizer(n_bits):
    class qfn(torch.autograd.Function):
        @staticmethod
        def forward(ctx, in_tensor):
            if n_bits == 32:
                # do not apply quantizer for full-bitwidth values
                out_tensor = in_tensor
            elif n_bits == 1:
                # Binary Weight Net implementation
                out_tensor = torch.sign(in_tensor)
            else:
                n = float(2 ** n_bits - 1)
                out = torch.round(in_tensor * n) / n
            return out

        @staticmethod
        def backward(ctx, grad_output):
            grad_input = grad_output.clone()
            return grad_input
    return qfn().apply<jupyter_output><empty_output><jupyter_text>Then we will need separate quantization functions for weights and activations.For weights the following formula is implemented:$$r_0 = 2 \cdot\textit{quantize_k}(\frac{tanh(r)}{2 * max|tanh(r)|} + \frac{1}{2}) - 1$$For activations uniform quantization is applied, as defined above.<jupyter_code>class weight_quantize_fn(nn.Module):
    def __init__(self, n_bits):
        super(weight_quantize_fn, self).__init__()
        assert n_bits <= 8 or n_bits == 32
        self.n_bits = n_bits
        self.value_range = 2 ** self.n_bits - 1
        self.uniform_q = get_uniform_quantizer(n_bits)

    def forward(self, x):
        if self.n_bits == 32:
            weight_q = x
        elif self.n_bits == 1:
            E = x.abs().mean().detach()
            weight_q = self.uniform_q(x / E) * E
        else:
            weight = x.tanh()
            weight = weight / 2 / weight.abs().max() + 0.5
            # weight_q = 2 * self.uniform_q(weight) - 1 # for original implementation with real values
            weight_q = self.uniform_q(weight)
            weight_q = self.value_range * weight_q - (self.value_range + 1) / 2
        return weight_q


class activation_quantize_fn(nn.Module):
    def __init__(self, n_bits):
        super(activation_quantize_fn, self).__init__()
        assert n_bits <= 8 or n_bits == 32
        self.n_bits = n_bits
        self.uniform_q = get_uniform_quantizer(n_bits)
    
    def forward(self, x):
        if self.n_bits == 32:
            activation_q = x
        else:
            activation_q = (2 ** self.n_bits - 1) * self.uniform_q(torch.clamp(x, 0, 1))
        return activation_q<jupyter_output><empty_output><jupyter_text>And now our building blocks allow us to implement basic convolution and activation functions.<jupyter_code>class DRFConv2d(nn.Conv2d):
    def __init__(self, n_bits, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=False):
        super(DRFConv2d, self).__init__(in_channels, out_channels, kernel_size, stride,
                                 padding, dilation, groups, bias)
        self.n_bits = n_bits
        self.quantize_fn = weight_quantize_fn(n_bits)

    def forward(self, in_tensor, order=None):
        weight_q = self.quantize_fn(self.weight)
        return F.conv2d(in_tensor, weight_q, self.bias, self.stride,
                  self.padding, self.dilation, self.groups)

class DRFActivation(nn.Module):
    def __init__(self, n_bits):
        super(DRFActivation, self).__init__()
        self.n_bits = n_bits
        self.quantize_fn = activation_quantize_fn(n_bits)

    def forward(self, in_tensor):
        if self.n_bits >= 32:
            return in_tensor
        out_tensor = self.quantize_fn(in_tensor)
        return out_tensor<jupyter_output><empty_output><jupyter_text>Let's implement the quantized CNN the same as the full-precision CNN<jupyter_code>class QNet(nn.Module):
    def __init__(self, n_bits=8):
        super(QNet, self).__init__()
        self.n_bits = n_bits
        self.conv1 = nn.Sequential(
            DRFConv2d(self.n_bits, 1, 32, 3, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            DRFActivation(self.n_bits)
        )
        self.maxpool_2x2 = nn.MaxPool2d((2,2))
        self.conv2 = nn.Sequential(
            DRFConv2d(self.n_bits, 32, 64, 3, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            DRFActivation(self.n_bits)
        )
        self.fc1 = DRFConv2d(self.n_bits, 64, 10, 1)
        self.gmp = nn.AdaptiveMaxPool2d(1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool_2x2(x)
        x = self.conv2(x)
        x = self.gmp(x)
        x = self.fc1(x)
        x = torch.flatten(x, 1)
        output = x / 100000
        output = F.log_softmax(output, dim=1)
        return output
# sanity check
qmodel = QNet().to(torch.device('cpu'))
in_tensor = torch.floor(torch.from_numpy(np.random.randint(0,255,size=(28,28))).view(1, 1, 28, 28).type(torch.float)) - 127
res = qmodel(in_tensor)
print(res.detach().numpy()[0])
model = QNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
scheduler = StepLR(optimizer, step_size=1, gamma=0.9)
for epoch in range(1, epochs + 1):
    train(model, device, train_loader, optimizer, epoch, log_interval=100)
    test(model, device, test_loader)
    scheduler.step()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1153122648607976, 'fraction_numerical': 0.02247930775018811, 'mean_word_length': 2.8317117117117117, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 37, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '315030', 'license': 'MIT', 'n_tokens_mistral': 3618, 'n_tokens_neox': 3416, 'n_words': 820, 'path': 'week1_14_segmentation_model_compression/week14_quantization_practice_additional.ipynb', 'repo_name': 'GendalfSeriyy/ml-mipt', 'size': '10632'}