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

username: Anmol_Gautam
cooked: <p>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.</p>
<p>Now when I run it, its showing error in custom conv2d and Linear layers as follows → <br>
<strong>" RuntimeError: Error(s) in loading state_dict for DataParallel:</strong><br>
**	Missing key(s) in state_dict: “module.conv1.weight”, “module.conv1.bias”. **<br>
**	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”. "**</p>
<p>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.</p>
<p>I am sharing my code below. —&gt;</p>
<blockquote>
<p>import torch<br>
import torch.nn as nn<br>
import torch.nn.parallel<br>
import torch.optim<br>
import torch.utils.data<br>
import torch.nn.functional as F<br>
from torch.autograd import Variable<br>
import distiller.quantization as quantization<br>
from distiller.apputils.checkpoint import load_checkpoint<br>
import torchvision<br>
import torchvision.transforms as transforms<br>
import numpy<br>
from bitstring import BitArray</p>
<p>import math<br>
from torch.nn.parameter import Parameter<br>
from torch.nn.modules.utils import _pair<br>
import torch.autograd.function as Function</p>
<p>######### torch.device(‘cpu’)<br>
########################################  Pre-Processing Cifar ##########################</p>
<p>transform = transforms.Compose([transforms.ToTensor(),   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])</p>
<p>trainset = torchvision.datasets.CIFAR10(root=‘./data’, train=True, download=True, transform=transform)<br>
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)</p>
<p>testset = torchvision.datasets.CIFAR10(root=‘./data’, train=False, download=True, transform=transform)<br>
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)</p>
<p>classes = (‘plane’, ‘car’, ‘bird’, ‘cat’, ‘deer’, ‘dog’, ‘frog’, ‘horse’, ‘ship’, ‘truck’)</p>
<p>##############################  Test 2 ################################<br>
#########################################################################################################</p>
<p>########################################  LINEAR UNROLLED  #################################</p>
<p>class LinearFunction(torch.autograd.Function):<br>
<span class="mention">@staticmethod</span><br>
def forward(ctx, input, weight, bias=None):<br>
ctx.save_for_backward(input, weight, bias)<br>
output = input.mm(weight.t())<br>
if bias is not None:<br>
output += bias.unsqueeze(0).expand_as(output)<br>
return output<br>
<span class="mention">@staticmethod</span><br>
def backward(ctx, grad_output):<br>
input, weight, bias = ctx.saved_variables<br>
grad_input = grad_weight = grad_bias = None<br>
if ctx.needs_input_grad[0]:<br>
grad_input = grad_output.mm(weight)<br>
if ctx.needs_input_grad[1]:<br>
grad_weight = grad_output.t().mm(input)<br>
if bias is not None and ctx.needs_input_grad[2]:<br>
grad_bias = grad_output.sum(0).squeeze(0)<br>
return grad_input, grad_weight, grad_bias</p>
<p>class customLinear(nn.Module):<br>
def <strong>init</strong>(self, input_features, output_features, bias=True):<br>
super(customLinear, self).<strong>init</strong>()<br>
self.input_features = input_features<br>
self.output_features = output_features<br>
self.weight = nn.Parameter(torch.Tensor(output_features, input_features))<br>
if bias:<br>
self.bias = nn.Parameter(torch.Tensor(output_features))<br>
else:<br>
self.register_parameter(‘bias’, None)<br>
self.weight.data.uniform_(-0.1, 0.1)<br>
if bias is not None:<br>
self.bias.data.uniform_(-0.1, 0.1)<br>
def forward(self, input):<br>
return LinearFunction.apply(input, self.weight, self.bias)</p>
<p>############################################   CONV2D - UNROLLED  ################################################################</p>
<p>class customConv(nn.Module):<br>
def <strong>init</strong>(self, n_channels, out_channels, kernel_size, dilation=1, padding=0, stride=1,bias=True):<br>
super(customConv, self).<strong>init</strong>()<br>
self.kernel_size = <em>pair(kernel_size)<br>
self.out_channels = out_channels<br>
self.dilation = <em>pair(dilation)<br>
self.padding = <em>pair(padding)<br>
self.stride = <em>pair(stride)<br>
self.n_channels = n_channels<br>
self.weight = Parameter(torch.Tensor(self.out_channels, self.out_channels, self.n_channels, self.kernel_size[0], self.kernel_size[1]))<br>
if bias:<br>
self.bias = Parameter(torch.Tensor(self.out_channels))<br>
self.flag = True<br>
else:<br>
self.register_parameter(‘bias’, None)<br>
self.flag = False<br>
self.reset_parameters()<br>
def reset_parameters(self):<br>
n = self.n_channels<br>
for k in self.kernel_size:<br>
n *= k<br>
stdv = 1. / math.sqrt(n)<br>
self.weight.data.uniform</em>(-stdv, stdv)<br>
if self.bias is not None:<br>
self.bias.data.uniform</em>(-stdv, stdv)<br>
def forward(self, input</em>):<br>
hout = ((input</em>.shape[2] + 2 * self.padding[0] - self.dilation[0] * (self.kernel_size[0]-1)-1)<a href="//self.stride" rel="noopener nofollow ugc">//self.stride</a>[0])+1<br>
wout = ((input_.shape[3] + 2 * self.padding[1] - self.dilation[1] * (self.kernel_size[1]-1)-1)<a href="//self.stride" rel="noopener nofollow ugc">//self.stride</a>[1])+1<br>
inputUnfolded = F.unfold(input_, kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride)<br>
weightunfold = F.unfold(self.weight[0], kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride)<br>
input_kernel_shape = inputUnfolded.size()[1]<br>
inputUnfolded = inputUnfolded.reshape(-1, input_kernel_shape)<br>
weightunfold = weightunfold.transpose(0, -1)<br>
weightunfold = weightunfold[0]<br>
if self.flag:<br>
convolvedOutput = torch.matmul(inputUnfolded, weightunfold) + self.bias<br>
else:<br>
convolvedOutput = torch.matmul(inputUnfolded, weightunfold)<br>
convolutionReconstruction = convolvedOutput.view(input_.shape[0], self.out_channels, hout, wout)<br>
return convolutionReconstruction</p>
<p>################################</p>
<p>class Simplenet(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Simplenet, self).<strong>init</strong>()<br>
self.conv1 = customConv(3, 6, 5)   ###  Using Custom Conv2d<br>
self.relu_conv1 = nn.ReLU()<br>
self.pool1 = nn.MaxPool2d(2, 2)<br>
self.conv2 = customConv(6, 16, 5)<br>
self.relu_conv2 = nn.ReLU()<br>
self.pool2 = nn.MaxPool2d(2, 2)<br>
self.fc1 = customLinear(16 * 5 * 5, 120)    ###  Using Custom Linear<br>
self.relu_fc1 = nn.ReLU()<br>
self.fc2 = customLinear(120, 84)<br>
self.relu_fc2 = nn.ReLU()<br>
self.fc3 = customLinear(84, 10)<br>
def forward(self, x):<br>
x = self.pool1(self.relu_conv1(self.conv1(x)))<br>
x = self.pool2(self.relu_conv2(self.conv2(x)))<br>
x = x.view(-1, 16 * 5 * 5)<br>
x = self.relu_fc1(self.fc1(x))<br>
x = self.relu_fc2(self.fc2(x))<br>
x = self.fc3(x)<br>
x = F.log_softmax(x, dim=1)<br>
return x</p>
<p>#########################################################################################################<br>
##################################   Loading Model and Quantization ###############################</p>
<p>model = Simplenet()<br>
model = torch.nn.DataParallel(model)</p>
<p>quantizer = quantization.PostTrainLinearQuantizer(model, model_activation_stats=‘/home/soumyadeep/distiller/examples/classifier_compression/logs/2019.10.10-234320/configs/acts_quantization_stats.yaml’)<br>
dummy_input = Variable(torch.randn(4, 3, 32, 32), requires_grad=True)<br>
quantizer.prepare_model(dummy_input)</p>
<p>checkpoint = torch.load(‘/home/soumyadeep/distiller/examples/classifier_compression/logs/2019.10.10-234320/quantized_checkpoint.pth.tar’)<br>
model.load_state_dict(checkpoint[‘state_dict’], strict= True)<br>
model = model.to(torch.device(‘cuda’))<br>
print(‘&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt; Model Loaded !!! &gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;&gt;’)</p>
<p>#################################################################################################<br>
######################################   Prediction Part ########################################</p>
<p>def test_label_predictions(model, device, test_loader):<br>
model.eval()<br>
actuals = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
predictions = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
correct = 0<br>
with torch.no_grad():<br>
for data, target in test_loader:<br>
data, target = data.to(device), target.to(device)<br>
output = model(data)<br>
prediction = output.argmax(dim=1, keepdim=True)<br>
actuals.extend(target.view_as(prediction))<br>
predictions.extend(prediction)<br>
correct += prediction.eq(target.view_as(prediction)).sum().item()<br>
return [i.item() for i in actuals], [i.item() for i in predictions], correct</p>
<p>#########################################################################################<br>
##########################################  Accuracy ####################################</p>
<p>device_name = ‘cuda’<br>
actuals, predictions, correct = test_label_predictions(model, device_name, testloader)<br>
######### print(“actuals —”,actuals)<br>
######### print(“prediction —”, predictions)<br>
total_num_data = len(testloader.dataset)<br>
accuracy = (float(correct / total_num_data)) * 100<br>
print(“–acuracy—”, accuracy)</p>
<p>##################################################################################################<br>
#########################################  CODE END ##############################################</p>
</blockquote>