FP16 and use_amp support?

[FlotingDream]

FP16 and use_amp support

[FlotingDream]

it can be down by this

from torch.cuda.amp import custom_fwd,custom_bwd
# gradients in the backward are received in the order of tensor as they were output in forward function
class PiXOperator(torch.autograd.Function):
   @staticmethod
   @custom_fwd(cast_inputs=torch.float32)
   def forward(ctx, zeta: int, tau: float,  input: torch.Tensor, p: torch.Tensor):
       outputs = pix_layer_cuda.forward(zeta, tau, input, p)
       ctx.save_for_backward(input, p)
       ctx.zeta = zeta
       ctx.tau = tau
       return outputs[0]

   @staticmethod
   @custom_bwd
   def backward(ctx, out_grad):
       input, p = ctx.saved_tensors
       zeta = ctx.zeta
       tau = ctx.tau
       input_grad, fusion_prob_grad = pix_layer_cuda.backward(zeta, tau, input, p, out_grad)
       return None, None, input_grad, fusion_prob_grad

but I want to know if the cuda'code can support fp16 operator? AND fp16 inference? even onnx suppot and trt support. like https://github.com/5had3z/CerberusNet/tree/master/runtime/cerberus_net/trt_plugins

please help, thx!

[ashishkumar822]

Thanks for the patience @FlotingDream . Currently, we have not written CUDA kernels for FP16. But the FP16 kernel can also be written.