Allow FX graph mode post-training dynamic quantisation of BlurConv2d operations.

[BrettRyland]

What does this PR do?

This PR extracts the dynamic control flow operations in blur_2d into static control flow operations in the BlurConv2d class while leaving the dynamic control flow operations in place for BlurMaxPool2d. This allows FX graph mode post-training dynamic quantisation to be applied to models where blurpool has been applied to convolutions (though not max pooling).

There still seems to be an issue with static quantisation however, which I'm unsure of how to fix:

RuntimeError: quantized::conv2d_prepack() is missing value for argument 'bias'. Declaration: quantized::conv2d_prepack(Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int groups) -> __torch__.torch.classes.quantized.Conv2dPackedParamsBase

What issue(s) does this change relate to?

This PR partially fixes issue #1985 in that blurpool applied to convolutions uses static control flow and no longer prevents FX graph mode post-training dynamic quantisation. Blurpool applied to max pool continues to use dynamic control flow as before.

Before submitting

• Have you read the contributor guidelines?
• [n] Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• [n/a] Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[nik-mosaic]

Stamping to unblock but will rely on @dblalock's review for correctness.

Let's please wait for Davis' review. I would like that before merging

[dskhudia]

@BrettRyland We will take a look at it and get back to you soon. Also will check the other quantization error.

[dblalock]

LGTM; basically just suggested a few local changes to polish it.

The one subtlety I found is that this would change the semantics when the filter spatial size is larger than the input + padding, depending on whether channels is supplied or inferred. So I suggested just removing the checks and having it always throw at the conv2d in this case.

The "suggestions" aren't super legible since Github won't let them span deletions, but the result is just this:

    if channels < 1:  # Use dynamic control flow
        _, channels, h, w = input.shape

        if (filter.shape[0] == 1) and (channels > 1):
            # assume filt is already a rank 4 tensor
            filter = filter.repeat((channels, 1, 1, 1))

    return F.conv2d(input, filter, stride=stride, padding=padding, groups=channels, bias=None)

[dskhudia]

@BrettRyland Your issue regarding RuntimeError: quantized::conv2d_prepack() seems like a bug in PyTorch fx tracing for quantization. A workaround for this is to not use keywords arguments for F.conv2d.

Instead of return F.conv2d(input, filter, stride=stride, padding=padding, groups=channels, bias=None) the following works:
from torch.nn.modules.utils import _pair
return F.conv2d(input, filter, None, _pair(stride), _pair(padding), _pair(1), channels)

Feel free to modify PR as suggested above.

[BrettRyland]

@BrettRyland Your issue regarding RuntimeError: quantized::conv2d_prepack() seems like a bug in PyTorch fx tracing for quantization. A workaround for this is to not use keywords arguments for F.conv2d.

Instead of return F.conv2d(input, filter, stride=stride, padding=padding, groups=channels, bias=None) the following works: from torch.nn.modules.utils import _pair return F.conv2d(input, filter, None, _pair(stride), _pair(padding), _pair(1), channels)

Feel free to modify PR as suggested above.

Thanks for this. I've updated the PR with all of the above suggestions other than the one I commented on above.

[dskhudia]

@BrettRyland It is indeed a PyTorch quantization issue. I created an issue for PyTorch with a minimal repro (pytorch/pytorch#95492)

import torch
import torch.nn as nn

from torch.ao.quantization import get_default_qconfig, QConfigMapping
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.nn import functional as F


class myConv2d_1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        return F.conv2d(input, self.weight, None, [1, 1], [0, 0], [1, 1], 1)

class myConv2d_2(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))

    def forward(self, input):
        return F.conv2d(input, self.weight, bias=None, stride=[1, 1], padding=[0, 0], dilation=[1, 1], groups=1)

# This works
#model = myConv2d_1(8, 64, 3)

# This doesn't work
model = myConv2d_2(8, 64, 3)

model.eval()
qconfig = get_default_qconfig("fbgemm")
qconfig_mapping = QConfigMapping().set_global(qconfig)

example_inputs = torch.randn(1, 8, 224, 224)

prepared_model = prepare_fx(model, qconfig_mapping, example_inputs)
model(example_inputs)
quantized_model = convert_fx(prepared_model)
out_q = quantized_model(example_inputs)
quantized_model.print_readable()

[dskhudia]

@dblalock will handle the suggestions in a followup PR. Merging this for the release.

[dskhudia]

@dblalock : Just checked all your comments in detail and I don't think anything else is needed for this PR. Please let me know if you think otherwise.