fix CONVBN & CONVBNRELU for resnet

mzh/resnet.py

@@ -207,6 +207,9 @@ class ResNet(nn.Module):
         
         super(ResNet, self).__init__()
 
+        self.mode = mode
+        self.n_exp = n_exp
+
         self.inplanes = 16 # 因为 CIFAR-10 图片较小，所以开始时需要更少的通道数
         GlobalVariables.SELF_INPLANES = self.inplanes
         print('resnet init:'+ str(GlobalVariables.SELF_INPLANES))
@@ -283,10 +286,19 @@ class ResNet(nn.Module):
         x = x.view(x.size(0), -1)
         x = self.fc(x)
 
-        return x
+        out = F.softmax(x,dim = 1)         # 这里不softmax也行 影响不大
+
+        return out
 
     def quantize(self, num_bits=8):
-      pass
+        self.qconvbnrelu1 = QConvBNReLU(self.conv1,self.bn1,qi=True,qo=True,num_bits=num_bits,n_exp=self.n_exp, mode=self.mode)
+        # 没有输入num_bits 需修改
+        self.layer1.quantize(num_bits=num_bits)  
+        self.layer2.quantize(num_bits=num_bits)
+        self.layer3.quantize(num_bits=num_bits)
+        self.layer4.quantize(num_bits=num_bits)
+        self.qavgpool1 = QAdaptiveAvgPool2d(qi=False,qo=True,num_bits=num_bits,n_exp=self.n_exp, mode=self.mode)
+        self.qfc1 = QLinear(self.fc,qi=False,qo=True,num_bits=num_bits,n_exp=self.n_exp, mode=self.mode)
 
     def quantize_forward(self, x):
         # for _, layer in self.quantize_layers.items():
@@ -294,24 +306,56 @@ class ResNet(nn.Module):
 
         # out = F.softmax(x, dim=1)
         # return out
-        pass
+        x = self.qconvbnrelu1(x)
+        x = self.layer1.quantize_forward(x)
+        x = self.layer2.quantize_forward(x)
+        x = self.layer3.quantize_forward(x)
+        x = self.layer4.quantize_forward(x)
+        x = self.qavgpool1(x)
+        x = x.view(x.size(0), -1)   
+        x = self.qfc1(x)
+        
+        out = F.softmax(x,dim = 1)         # 这里不softmax也行 影响不大
+        return out
       
 
     def freeze(self):
-        pass
+        self.qconvbnrelu1.freeze()  # 因为作为第一层是有qi的，所以freeze的时候无需再重新提供qi
+        qo = self.layer1.freeze(qinput = self.qconvbnrelu1.qo)
+        qo = self.layer2.freeze(qinput = qo)
+        qo = self.layer3.freeze(qinput = qo)
+        qo = self.layer4.freeze(qinput = qo)
+        self.qavgpool1.freeze(qo)
+        self.qfc1.freeze(qi=qo)
 
     def fakefreeze(self):
         pass
 
     def quantize_inference(self, x):
-        pass
+        qx = self.qconvbnrelu1.qi.quantize_tensor(x,mode=self.mode)
+        qx = self.qconvbnrelu1.quantize_inference(qx)
+        qx = self.layer1.quantize_inference(qx)
+        qx = self.layer2.quantize_inference(qx)
+        qx = self.layer3.quantize_inference(qx)
+        qx = self.layer4.quantize_inference(qx)
+        qx = self.qavgpool1.quantize_inference(qx)
+        qx = qx.view(qx.size(0), -1)
+        qx = self.qfc1.quantize_inference(qx) 
+        if self.mode == 1:
+            qx = self.qfc1.qo.dequantize_tensor(qx,mode=self.mode)
+        
+
+        out = F.softmax(qx,dim = 1)         # 这里不softmax也行 影响不大
+        return out
+
+
 
 
 # BasicBlock 类
 class BasicBlock(nn.Module):
     expansion = 1
 
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
+    def __init__(self, inplanes, planes, stride=1, downsample=None,  n_exp=4, mode=1):
         super(BasicBlock, self).__init__()
 
         # 第一个卷积层
@@ -328,6 +372,8 @@ class BasicBlock(nn.Module):
         self.relu = nn.ReLU()
         self.downsample = downsample
         self.stride = stride
+        self.mode = mode
+        self.n_exp = n_exp
 
     def forward(self, x):
 
@@ -349,13 +395,13 @@ class BasicBlock(nn.Module):
         return out
     
     def quantize(self, num_bits=8):
-        self.qconvbnrelu1 = QConvBNReLU(self.conv1,self.bn2,qi=False,qo=True,num_bits=num_bits)
-        self.qconvbn1 = QConvBN(self.conv2,self.bn2,qi=False,qo=True,num_bits=num_bits)
+        self.qconvbnrelu1 = QConvBNReLU(self.conv1,self.bn1,qi=False,qo=True,num_bits=num_bits,n_exp=self.n_exp,mode=self.mode)
+        self.qconvbn1 = QConvBN(self.conv2,self.bn2,qi=False,qo=True,num_bits=num_bits,n_exp=self.n_exp,mode=self.mode)
 
         if self.downsample is not None:
-            self.qconvbn2 =  QConvBN(self.downsample[0],self.downsample[1],qi=False,qo=True,num_bits=num_bits)
+            self.qconvbn2 =  QConvBN(self.downsample[0],self.downsample[1],qi=False,qo=True,num_bits=num_bits,n_exp=self.n_exp,mode=self.mode)
 
-        self.qrelu1 = QReLU()
+        self.qrelu1 = QReLU(n_exp=self.n_exp,mode=self.mode)
         
 
     def quantize_forward(self, x):
@@ -371,24 +417,26 @@ class BasicBlock(nn.Module):
         out = self.qrelu1(out)
         return out
 
-    def freeze(self):
+    def freeze(self, qinput):
         # 这里的qconvbnrelu1其实是可以用前一层的qo的，但感觉不太好传参，就没用 
         # 还需仔细检查
-        self.qconvbnrelu1.freeze()
+        self.qconvbnrelu1.freeze(qi= qinput)   # 需要接前一个module的最后一个qo
         self.qconvbn1.freeze(qi = self.qconvbnrelu1.qo)
 
         if self.downsample is not None:
-            self.qconvbn2.freeze(qi = self.qconvbn1)
-            self.qrelu1.freeze(self.qconvbn2)
+            self.qconvbn2.freeze(qi = self.qconvbn1.qo)
+            self.qrelu1.freeze(self.qconvbn2.qo)
+            return self.qconvbn2.qo
 
         else:
-            self.qrelu1.freeze(self.qconvbn1)
+            self.qrelu1.freeze(self.qconvbn1.qo)
+            return self.qconvbn1.qo
 
     def quantize_inference(self, x):
         # 感觉是不需要进行初始的quantize_tensor和dequantize_tensor，因为他不是最前/后一层，只要中间的每层都在量化后的领域内，就不需要这种处理。
         identity = x
         out = self.qconvbnrelu1.quantize_inference(x)
-        out = self.qconvbn1.quantize_inference(x)
+        out = self.qconvbn1.quantize_inference(out)
 
         if self.downsample is not None:
             identity = self.qconvbn2.quantize_inference(identity)
@@ -453,7 +501,7 @@ class Bottleneck(nn.Module):
 
 class MakeLayer(nn.Module):
 
-    def __init__(self, block, planes, blocks, stride=1):
+    def __init__(self, block, planes, blocks, stride=1, n_exp=4, mode=1):
         super(MakeLayer, self).__init__()
         print('makelayer init:'+ str(GlobalVariables.SELF_INPLANES))
         self.downsample = None
@@ -462,12 +510,15 @@ class MakeLayer(nn.Module):
             nn.Conv2d(GlobalVariables.SELF_INPLANES, planes * block.expansion,kernel_size=1, stride=stride, bias=False), 
             nn.BatchNorm2d(planes * block.expansion)
             )
-        
+        self.n_exp = n_exp
+        self.mode = mode
         self.blockdict = nn.ModuleDict()
-        self.blockdict['block1'] = block(GlobalVariables.SELF_INPLANES, planes, stride, self.downsample)
+        self.blockdict['block1'] = block(inplanes=GlobalVariables.SELF_INPLANES, planes=planes, stride=stride, downsample=self.downsample,n_exp=self.n_exp,mode=self.mode)
         GlobalVariables.SELF_INPLANES = planes * block.expansion
         for i in range(1, blocks):  # block的个数   这里只能用字典了
-            self.blockdict['block' + str(i+1)] = block(GlobalVariables.SELF_INPLANES, planes)  # 此处进行实例化了
+            self.blockdict['block' + str(i+1)] = block(inplanes=GlobalVariables.SELF_INPLANES, planes=planes,n_exp=self.n_exp, mode=self.mode)  # 此处进行实例化了
+        
+        
 
     # def _make_layer(self, block, planes, blocks, stride=1):
     #     downsample = None
@@ -499,7 +550,7 @@ class MakeLayer(nn.Module):
     def quantize(self, num_bits=8):
         # 需检查
         for _, layer in self.blockdict.items():
-            layer.quantize()   # 这里是因为每一块都是block，而block中有具体的quantize策略
+            layer.quantize(num_bits=num_bits)   # 这里是因为每一块都是block，而block中有具体的quantize策略, n_exp和mode已经在__init__中赋值了
         
 
     def quantize_forward(self, x):
@@ -509,11 +560,18 @@ class MakeLayer(nn.Module):
         return x
         
        
-    def freeze(self):
+    def freeze(self, qinput):  # 需要在 Module Resnet的freeze里传出来
         # 这里的qconvbnrelu1其实是可以用前一层的qo的，但感觉不太好传参，就没用 
         # 还需仔细检查
+        cnt = 0
         for _, layer in self.blockdict.items():
-            layer.freeze()  # 各个block中有具体的freeze
+            if cnt == 0:
+                qo = layer.freeze(qinput = qinput)
+                cnt = 1
+            else:
+                qo = layer.freeze(qinput = qo)  # 各个block中有具体的freeze
+
+        return qo   # 供后续的层用
 
 
     def quantize_inference(self, x):