Multi-pretrained weight support - Quantized ResNet50

[datumbox]

Resolves #4671

Example usage:

from PIL import Image
from torchvision import prototype as P


img = Image.open("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

combos = [
    (P.models.quantization.QuantizedResNet50Weights.ImageNet1K_FBGEMM_RefV1, True),
    (P.models.ResNet50Weights.ImageNet1K_RefV1, False)
]

for weights, quantize in combos:
    # Initialize model
    print(weights, quantize)
    model = P.models.quantization.resnet50(weights=weights, quantize=quantize)
    model.eval()

    # Initialize inference transforms
    preprocess = weights.transforms()

    # Apply inference preprocessing transforms
    batch = preprocess(img).unsqueeze(0)
    prediction = model(batch).squeeze(0).softmax(0)

    # Make predictions
    label = prediction.argmax().item()
    score = prediction[label].item()

    # Use meta to get label
    category_name = weights.meta['categories'][label]
    print(f"{category_name}: {100 * score}%")

cc @pmeier @bjuncek

[datumbox]

Providing some clarifications on comments:

[datumbox]

Fixing in place a previous "bug" on our typing. The correct types here are the more specific Quantizable versions.

[datumbox]

Not all weights are expected to have the "backend" meta. For example when quantize=False is passed.

Let's ignore the silent overwriting of parameters (see #4613 (comment)). This can be discussed separately and applied everywhere on a follow up.

[datumbox]

The backend is allowed to be a "hidden" kwargs argument. We can decide making it public on future iterations.

[datumbox]

We allow passing both Quantized and normal weights. This is aligned with the past behaviours where different URLs were loaded depending on the value of quantize.

[datumbox]

The value of which default weights we load depends on whether we quantize or not.

[datumbox]

The validation of the weights also depends on the value of quantize. Passing the wrong combination will throw an error.

[NicolasHug]

stamping

[prabhat00155]

Thanks @datumbox!

[fmassa]

Thanks for the PR!

I'm approving to unblock. The current structure follows somewhat closely what we do in the main branch, so it's a good start. I've added some comments with some potential ideas to consider in the future.

Also, it would be good to get someone from the quantization team to have a look and see if the weight structure (with backends etc) is still relevant as of today, or if there will be new things coming that we should consider (sparsity, etc).

[fmassa]

Flagging again #4613 (comment) but we can discuss in a follow-up

[fmassa]

Throwing an idea in the wild: as of today (and I think this will be the case for all models), all quantized weights originates from an unquantized weight. Do we want to keep this link somehow in the Weights structure? Do we want the quantized weights to be magically picked from the ResNet50Weights if we pass the quantize flag?

I think it might be important to keep this relationship somehow.

[datumbox]

We could keep this in the meta-data or if we think we are willing to make a commitment pass this as a proper field of the Weights Data Class. The reason we don't is because quantization is not something too mature at the moment. Only classification models have been quantized and we examine alternative APIs/approaches (such as FX) to achieve it. For these reasons, I would be in favour of not introducing a direct link and review this decision on the near future.

[fmassa]

nit: you could probably simplify a tiny bit the code by doing something like

if quantize:
    weights_table = QuantizedResNet50Weights
else:
    weights_table = ResNet50Weights
...
weights = weights_table.ImageNet1K_RefV1  # different naming convention than now
weights = weights_table.verify(weights)

In some sense, it's a bit annoying to have to carry those two Weights inside every model builder

[datumbox]

I decided not to connect the two using the same name because the Reference/recipe version might not be the same; aka RefV1 for quantized is something different than for non-quantized. For example let's say you use a different recipe / config to achieve even better quantization of the same weights; now you need a different version of the enum. Moreover there might be multiple quantized weights enums for the same unquantized weights (for example if multiple backends should be supported). Some of these points will become clearer on the near future and we can revisit them.