Custom Conv2d and Custom Linear model creation issue with nn.Dataparallel

Firstly, I trained the cifar-10 dataset with distiller simple-prenet and quantized it to activation and weight of both 8bits using distiller again. After that, I am designed my own Custom Linear and Conv2d layers so that I can unroll it to access the matmul function using the quantized pretrained model.

Now when I run it, its showing error in custom conv2d and Linear layers as follows →
" RuntimeError: Error(s) in loading state_dict for DataParallel:
** Missing key(s) in state_dict: “module.conv1.weight”, “module.conv1.bias”. **
** Unexpected key(s) in state_dict: “module.conv1.in_0_scale”, “module.conv1.in_0_zero_point”, “module.conv1.output_scale”, “module.conv1.output_zero_point”, “module.conv1.w_scale”, “module.conv1.w_zero_point”, “module.conv1.accum_scale”, “module.conv1.is_simulated_quant_weight_shifted”, “module.conv1.wrapped_module.weight”, “module.conv1.wrapped_module.bias”. "**

Then I changed, load_state_dict, strict = False which resulted in running but the accuracy dropped to 8.86% which is abnormal. I tried a lot to debug this but couldn’t find the answer. Can you please check my code and write or tell me the correct layers or code for me.

I am sharing my code below. —>

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.nn.functional as F
from torch.autograd import Variable
import distiller.quantization as quantization
from distiller.apputils.checkpoint import load_checkpoint
import torchvision
import torchvision.transforms as transforms
import numpy
from bitstring import BitArray

import math
from torch.nn.parameter import Parameter
from torch.nn.modules.utils import _pair
import torch.autograd.function as Function

######### torch.device(‘cpu’)
######################################## Pre-Processing Cifar ##########################

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root=‘./data’, train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root=‘./data’, train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

classes = (‘plane’, ‘car’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)

############################## Test 2 ################################
#########################################################################################################

######################################## LINEAR UNROLLED #################################

class LinearFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, bias=None):
ctx.save_for_backward(input, weight, bias)
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_variables
grad_input = grad_weight = grad_bias = None
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0).squeeze(0)
return grad_input, grad_weight, grad_bias

class customLinear(nn.Module):
def init(self, input_features, output_features, bias=True):
super(customLinear, self).init()
self.input_features = input_features
self.output_features = output_features
self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(output_features))
else:
self.register_parameter(‘bias’, None)
self.weight.data.uniform_(-0.1, 0.1)
if bias is not None:
self.bias.data.uniform_(-0.1, 0.1)
def forward(self, input):
return LinearFunction.apply(input, self.weight, self.bias)

############################################ CONV2D - UNROLLED ################################################################

class customConv(nn.Module):
def init(self, n_channels, out_channels, kernel_size, dilation=1, padding=0, stride=1,bias=True):
super(customConv, self).init()
self.kernel_size = pair(kernel_size)
self.out_channels = out_channels
self.dilation = pair(dilation)
self.padding = pair(padding)
self.stride = pair(stride)
self.n_channels = n_channels
self.weight = Parameter(torch.Tensor(self.out_channels, self.out_channels, self.n_channels, self.kernel_size[0], self.kernel_size[1]))
if bias:
self.bias = Parameter(torch.Tensor(self.out_channels))
self.flag = True
else:
self.register_parameter(‘bias’, None)
self.flag = False
self.reset_parameters()
def reset_parameters(self):
n = self.n_channels
for k in self.kernel_size:
n *= k
stdv = 1. / math.sqrt(n)
self.weight.data.uniform
(-stdv, stdv)
if self.bias is not None:
self.bias.data.uniform
(-stdv, stdv)
def forward(self, input
):
hout = ((input
.shape[2] + 2 * self.padding[0] - self.dilation[0] * (self.kernel_size[0]-1)-1)//self.stride[0])+1
wout = ((input_.shape[3] + 2 * self.padding[1] - self.dilation[1] * (self.kernel_size[1]-1)-1)//self.stride[1])+1
inputUnfolded = F.unfold(input_, kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride)
weightunfold = F.unfold(self.weight[0], kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride)
input_kernel_shape = inputUnfolded.size()[1]
inputUnfolded = inputUnfolded.reshape(-1, input_kernel_shape)
weightunfold = weightunfold.transpose(0, -1)
weightunfold = weightunfold[0]
if self.flag:
convolvedOutput = torch.matmul(inputUnfolded, weightunfold) + self.bias
else:
convolvedOutput = torch.matmul(inputUnfolded, weightunfold)
convolutionReconstruction = convolvedOutput.view(input_.shape[0], self.out_channels, hout, wout)
return convolutionReconstruction

################################

class Simplenet(nn.Module):
def init(self):
super(Simplenet, self).init()
self.conv1 = customConv(3, 6, 5) ### Using Custom Conv2d
self.relu_conv1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = customConv(6, 16, 5)
self.relu_conv2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = customLinear(16 * 5 * 5, 120) ### Using Custom Linear
self.relu_fc1 = nn.ReLU()
self.fc2 = customLinear(120, 84)
self.relu_fc2 = nn.ReLU()
self.fc3 = customLinear(84, 10)
def forward(self, x):
x = self.pool1(self.relu_conv1(self.conv1(x)))
x = self.pool2(self.relu_conv2(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = self.relu_fc1(self.fc1(x))
x = self.relu_fc2(self.fc2(x))
x = self.fc3(x)
x = F.log_softmax(x, dim=1)
return x

#########################################################################################################
################################## Loading Model and Quantization ###############################

model = Simplenet()
model = torch.nn.DataParallel(model)

quantizer = quantization.PostTrainLinearQuantizer(model, model_activation_stats=‘/home/soumyadeep/distiller/examples/classifier_compression/logs/2019.10.10-234320/configs/acts_quantization_stats.yaml’)
dummy_input = Variable(torch.randn(4, 3, 32, 32), requires_grad=True)
quantizer.prepare_model(dummy_input)

checkpoint = torch.load(‘/home/soumyadeep/distiller/examples/classifier_compression/logs/2019.10.10-234320/quantized_checkpoint.pth.tar’)
model.load_state_dict(checkpoint[‘state_dict’], strict= True)
model = model.to(torch.device(‘cuda’))
print(‘<<<<<<<<<<<<<<<<< Model Loaded !!! >>>>>>>>>>>>>>>>>>’)

#################################################################################################
###################################### Prediction Part ########################################

def test_label_predictions(model, device, test_loader):
model.eval()
actuals =
predictions =
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
prediction = output.argmax(dim=1, keepdim=True)
actuals.extend(target.view_as(prediction))
predictions.extend(prediction)
correct += prediction.eq(target.view_as(prediction)).sum().item()
return [i.item() for i in actuals], [i.item() for i in predictions], correct

#########################################################################################
########################################## Accuracy ####################################

device_name = ‘cuda’
actuals, predictions, correct = test_label_predictions(model, device_name, testloader)
######### print(“actuals —”,actuals)
######### print(“prediction —”, predictions)
total_num_data = len(testloader.dataset)
accuracy = (float(correct / total_num_data)) * 100
print(“–acuracy—”, accuracy)

##################################################################################################
######################################### CODE END ##############################################