How can I incorporate PReLU in a quantized model?

Shisho_Sama: <p>Hello everyone.<br>
This is a followup question concerning <a href="https://discuss.pytorch.org/t/runtimeerror-could-not-run-aten-native-batch-norm-with-arguments-from-the-quantizedcputensorid-backend-aten-native-batch-norm-is-only-available-for-these-backends-cputensorid-mkldnncputensorid-variabletensorid/88975">this</a> . The issue is in the Resnet model that I’m dealing with, I cant replace <code>PReLU</code> with <code>ReLU</code> as it drastically affects the network performance.<br>
So my question is, what are my options here? what should I be doing in this case?<br>
Would doing sth like this suffice?</p>
<pre><code class="lang-python">class PReLU_Quantized(nn.Module):
    def __init__(self, prelu_object):
        super().__init__()
        self.weight = prelu_object.weight
        self.quantized_op = nn.quantized.FloatFunctional()
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, inputs):
        # inputs = torch.max(0, inputs) + self.weight * torch.min(0, inputs)    
        self.weight = self.quant(self.weight)
        weight_min_res = self.quantized_op.mul(self.weight, torch.min(inputs)[0])
        inputs = self.quantized_op.add(torch.max(inputs)[0], weight_min_res).unsqueeze(0)
        self.weight = self.dequant(self.weight)
        return inputs
</code></pre>
<p>and for the replacement :</p>
<pre><code class="lang-auto">class model(nn.Module):
     def __init__(self)
         super().__init__()
         .... 
        self.prelu = PReLU()
        self.prelu_q = PReLU_Quantized(self.prelu)
         ....
</code></pre>
<p>Thanks a lot in advance</p>

Shisho_Sama: <p>for some reason, the error between the actual PReLU and my implementation is  very large!<br>
here are sample diffs in different layers:</p>
<pre><code class="lang-python">diff : 1.1562038660049438
diff : 0.02868632599711418
diff : 0.3653906583786011
diff : 1.6100226640701294
diff : 0.8999372720718384
diff : 0.03773299604654312
diff : -0.5090572834014893
diff : 0.1654307246208191
diff : 1.161868691444397
diff : 0.026089997962117195
diff : 0.4205571115016937
diff : 1.5337920188903809
diff : 0.8799554705619812
diff : 0.03827812895178795
diff : -0.40296515822410583
diff : 0.15618863701820374
</code></pre>
<p>and the diff is calculated like this in the forward pass:</p>
<pre><code class="lang-python">    def forward(self, x):
        residual = x
        out = self.bn0(x)
        out = self.conv1(out)
        out = self.bn1(out)
        out = self.prelu(out)
        
        out2 = self.prelu2(out)
        print(f'diff : {( out - out2).mean().item()}')

        out = self.conv2(out)
</code></pre>
<p>This is the normal implementation which I used on ordinary model (i.e. not quantized!)  to assess whether it produces correct result and then move on to quantized version:</p>
<pre><code class="lang-python">class PReLU_2(nn.Module):
    def __init__(self, prelu_object):
        super().__init__()
        self.prelu_weight = prelu_object.weight
        self.weight = self.prelu_weight

    def forward(self, inputs):
        x = self.weight
        tmin, _ = torch.min(inputs,dim=0)
        tmax, _ = torch.max(inputs,dim=0)
        weight_min_res = torch.mul(x, tmin)
        inputs = torch.add(tmax, weight_min_res)
        inputs = inputs.unsqueeze(0)
        return inputs

</code></pre>
<p>what am I missing here?</p>

Shisho_Sama: <p>OK, I figured it out! I made a huge mistake in the very begining. I needed to calculate</p>
<pre><code class="lang-auto">PReLU(x)=max(0,x)+a∗min(0,x)
</code></pre>
<p>or<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/a/4/a4683d696a2a7178e866b53b58efa84054b22f4b.png" alt="image" data-base62-sha1="nspDyO6leATlzwvsEWOc3Qp1g15" width="337" height="88"><br>
and not the actual min! or max! which doesnt make sense!<br>
now, can anyone do me a favor and tell me how I can vectorize this ? I’m kind of lost at the moment!</p>

Shisho_Sama: <p>Thanks to dear God its done!<br>
Here is the final solution!:</p>
<pre><code class="lang-python">class PReLU_2(nn.Module):
    def __init__(self, prelu_object):
        super().__init__()
        self.prelu_weight = prelu_object.weight
        self.weight = self.prelu_weight

    def forward(self, inputs):
        pos = torch.relu(inputs)
        neg = -self.weight * torch.relu(-inputs)
        inputs = pos + neg
        return inputs
</code></pre>
<p>and t his is the quantized version :</p>
<pre><code class="lang-python">class PReLU_Quantized(nn.Module):
    def __init__(self, prelu_object):
        super().__init__()
        self.prelu_weight = prelu_object.weight
        self.weight = self.prelu_weight
        self.quantized_op = nn.quantized.FloatFunctional()
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, inputs):
        # inputs = max(0, inputs) + alpha * min(0, inputs) 
        self.weight = self.quant(self.weight)
        weight_min_res = self.quantized_op.mul(-self.weight, torch.relu(-inputs))
        inputs = self.quantized_op.add(torch.relu(inputs), weight_min_res)
        inputs = self.dequant(inputs)
        self.weight = self.dequant(self.weight)
        return inputs

</code></pre>

viash: <p>Hi, Thanks for the solution on PRELU Quantized. Any Tips on fusing the PRELU with CONV<br>
CONVPRELU , CONVBNPRELU.</p>