Hello,
I am trying to use your ttqConv2D class in a simple test case with mnist dataset.
This is the original structure:
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5, bias=False)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5, bias=False)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        # x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = F.relu(F.max_pool2d(self.conv2(x),2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

Was able to get to 99% accuracy after 10 epoch.

image.jpg849×289 143 KB

I then tried to replace nn.Conv2D with the following:
        super(Net, self).__init__()
        self.conv1 = ttqConv2D(1,10,5)
        self.conv2 = ttqConv2D(10,20,5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

but the model won’t converge in this case and will get stuck at ~ 11%. I then tried to modify your ttqConv2d class so that it actually doesn’t manipulate the weight nor gradient, but the model still doesn’t converge. Do you know what I did wrong?
Revised ttqConv2d without the quantization
class ternarizeWeights(torch.autograd.Function):

    def __init__(self, ternThres):
        super(ternarizeWeights, self).__init__()
        self.tThresScale = ternThres

    def forward(self, img2Conv, fpWeights, tWp, tWn):
        self.save_for_backward(img2Conv, fpWeights, tWp, tWn)
        # compute the ternarized weights with the given threshold

        quantW = fpWeights.clone()
        quantThres = self.tThresScale * quantW.abs().max()

        return img2Conv, quantW

    def backward(self, grad_outputImg, grad_outputW):
        img2Conv, fpWeights, tWp, tWn = self.saved_tensors
        quantThres = self.tThresScale * fpWeights.abs().max()
        # leave the image gradient unchanged
        grad_img2Conv = grad_outputImg
        # all other gradients have to be manipulated
        # first clone them, to ensure the correct datatype and dimension
        grad_fpWeights = fpWeights.clone()
        grad_tWp = tWp.clone()
        grad_tWn = tWn.clone()

        return grad_img2Conv, grad_fpWeights, grad_tWp, grad_tWn


# write the second autograd function -> override forward
# call the standard convolution function on the unchanged image but
# with ternarized weights (2D case) and pass back their gradients
# -> the predeceding layer has to handle the gradients

class execute2DConvolution(torch.nn.Module):
    def __init__(self, inStride=1, inPadding=0, inDilation=1, inGroups=1):
        super(execute2DConvolution, self).__init__()
        self.cStride = inStride
        self.cPad = inPadding
        self.cDil = inDilation
        self.cGrp = inGroups

    def forward(self, dataIn, weightIn):
        return torch.nn.functional.conv2d(dataIn, weightIn, bias=None,
                                          stride=self.cStride, padding=self.cPad,
                                          dilation=self.cDil, groups=self.cGrp)


# wrap the two functions inside one module, so this appears as a
# customized convolution to the outside world as Conv1/2D

class ttqConv2D(torch.nn.Module):
    def __init__(self, inWCin, inWCout, inWH,
                 inStride=1, inPad=0, inDil=1, inGroups=1, inTScale=0.05):
        super(ttqConv2D, self).__init__()
        # initialize all parameters that the convolution function needs to know
        self.conStride = inStride
        self.conPad = inPad
        self.outPad = 0
        self.conDil = inDil
        self.conTrans = False
        self.conGroups = inGroups

        # initialize the weights and the bias as well as the
        self.tScale = inTScale
        self.fpWeight = torch.nn.Parameter(torch.Tensor(inWCout, inWCin, inWH, inWH))
        # xavier weight initialization
        nnInit.xavier_normal(self.fpWeight)
        self.tWeightPos = torch.nn.Parameter(torch.Tensor([0]))
        self.tWeightNeg = torch.nn.Parameter(torch.Tensor([0]))
        nnInit.uniform(self.tWeightPos, 0.8, 1.2)
        nnInit.uniform(self.tWeightNeg, 0.8, 1.2)

    def forward(self, dataInput):
        dInput, tWeights = ternarizeWeights(self.tScale)(dataInput, self.fpWeight,
                                                         self.tWeightPos, self.tWeightNeg)

        return execute2DConvolution(self.conStride, self.conPad,
                                    self.conDil, self.conGroups)(dInput, tWeights)

Thanks,