How to quantize torch.nn.functional.grid_sample?

1

Hi all,
I want to convert the inputs of grid_sample from float32 to quint8 but it seems that the pytorch doesn’t support grid_sample in quint8.
How to quantize torch.nn.functional.grid_sample in quint8?
It’s there any other operation that can replace torch.nn.functional.grid_sample and can be quantized in quint8?
Below is the part of my code.
Thanks.

from re import T
import torch.nn.functional as F
import torch.nn as nn
import torch
from torch import Tensor

from torch.quantization import QuantStub, DeQuantStub
import torch.quantization

use_quan = True

class FlowWarp(nn.Module):
def init(self, interpolation=‘bilinear’, padding_mode=‘zeros’, align_corners=True):
super().init()
self.interpolation=interpolation
self.padding_mode=padding_mode
self.align_corners=align_corners

    if(use_quan):
        self.quant = QuantStub()
        self.quant_x = QuantStub()
        self.dequant = DeQuantStub()
        self.f = torch.nn.quantized.FloatFunctional()

def forward(self, x, flow: Tensor) -> Tensor:
    if flow.shape[-1] != 2:
        flow = flow.permute(0, 2, 3, 1)

    if(use_quan):
        x = self.dequant(x)
    n, c, h, w = x.size()

    grid_y, grid_x = torch.meshgrid(
        torch.linspace(-1, 1, h),
        torch.linspace(-1, 1, w),
        indexing='ij'
    )

    grid = torch.stack((grid_x, grid_y), 2).type_as(x).expand(n, h, w, 2)
    grid.requires_grad = False

    if(use_quan):
        grid = self.quant(grid)
        x = self.quant_x(x)
        grid_flow = self.f.add(grid, flow)
    else:
        grid_flow = grid + flow
    output = F.grid_sample(
        x,
        grid_flow,
        mode=self.interpolation,
        padding_mode=self.padding_mode,
        align_corners=self.align_corners
    )
    return output