breadlicker45/gpuGAN

1# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9"""Custom replacement for `torch.nn.functional.grid_sample` that10supports arbitrarily high order gradients between the input and output.11Only works on 2D images and assumes12`mode='bilinear'`, `padding_mode='zeros'`, `align_corners=False`."""13 14import warnings15import torch16 17# pylint: disable=redefined-builtin18# pylint: disable=arguments-differ19# pylint: disable=protected-access20 21#----------------------------------------------------------------------------22 23enabled = False  # Enable the custom op by setting this to true.24 25#----------------------------------------------------------------------------26 27def grid_sample(input, grid):28    if _should_use_custom_op():29        return _GridSample2dForward.apply(input, grid)30    return torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)31 32#----------------------------------------------------------------------------33 34def _should_use_custom_op():35    if not enabled:36        return False37    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):38        return True39    warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')40    return False41 42#----------------------------------------------------------------------------43 44class _GridSample2dForward(torch.autograd.Function):45    @staticmethod46    def forward(ctx, input, grid):47        assert input.ndim == 448        assert grid.ndim == 449        output = torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)50        ctx.save_for_backward(input, grid)51        return output52 53    @staticmethod54    def backward(ctx, grad_output):55        input, grid = ctx.saved_tensors56        grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid)57        return grad_input, grad_grid58 59#----------------------------------------------------------------------------60 61class _GridSample2dBackward(torch.autograd.Function):62    @staticmethod63    def forward(ctx, grad_output, input, grid):64        op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward')65        grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)66        ctx.save_for_backward(grid)67        return grad_input, grad_grid68 69    @staticmethod70    def backward(ctx, grad2_grad_input, grad2_grad_grid):71        _ = grad2_grad_grid # unused72        grid, = ctx.saved_tensors73        grad2_grad_output = None74        grad2_input = None75        grad2_grid = None76 77        if ctx.needs_input_grad[0]:78            grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid)79 80        assert not ctx.needs_input_grad[2]81        return grad2_grad_output, grad2_input, grad2_grid82 83#----------------------------------------------------------------------------84