Why nn.Upsample/F.interpolate followed by nn.InstanceNorm2d will report error "Unsupported: ONNX export of instance_norm for unknown channel size."
Labels: module: onnx, triaged, module: norms and normalization

### 🐛 Describe the bug

Hi~

When run **python main.py**, the error is

`torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of instance_norm for unknown channel size.  [Caused by the value 'input.199 defined in (%input.199 : Float(*, *, *, *, strides=[262144, 1024, 32, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[2, 2], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%input.195, %middle.1.weight, %middle.1.bias), scope: __main__.InpaintGenerator::/torch.nn.modules.container.Sequential::middle/torch.nn.modules.conv.Conv2d::middle.1`

Having tried:
1. Changing opset from 9 to 17
2. Different version of Pytorch (1.9, 1.10, 1.12, 1.13, 2.0)

But after removing IN layers, it works. So how to solve it? 
Thanks~

By the way, the code of the network architecture is from https://github.com/tsingqguo/misf. 


**Code File 1: ./main.py**

    import torch
    import torch.nn as nn
    from kpn.network import KernelConv
    import kpn.utils as kpn_utils
    import numpy as np
    
    
    class BaseNetwork(nn.Module):
        def __init__(self):
            super(BaseNetwork, self).__init__()
    
        def init_weights(self, init_type='normal', gain=0.02):
            def init_func(m):
                classname = m.__class__.__name__
                if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
                    if init_type == 'normal':
                        nn.init.normal_(m.weight.data, 0.0, gain)
                    elif init_type == 'xavier':
                        nn.init.xavier_normal_(m.weight.data, gain=gain)
                    elif init_type == 'kaiming':
                        nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                    elif init_type == 'orthogonal':
                        nn.init.orthogonal_(m.weight.data, gain=gain)
    
                    if hasattr(m, 'bias') and m.bias is not None:
                        nn.init.constant_(m.bias.data, 0.0)
    
                elif classname.find('BatchNorm2d') != -1:
                    nn.init.normal_(m.weight.data, 1.0, gain)
                    nn.init.constant_(m.bias.data, 0.0)
    
            self.apply(init_func)
    
    
    class InpaintGenerator(BaseNetwork):
        def __init__(self, config=None, residual_blocks=8, init_weights=True):
            super(InpaintGenerator, self).__init__()
    
            # self.filter_type = config.FILTER_TYPE
            # self.kernel_size = config.kernel_size
    
            self.encoder0 = nn.Sequential(
                nn.ReflectionPad2d(3),
                nn.Conv2d(in_channels=4, out_channels=64, kernel_size=7, padding=0),
                nn.InstanceNorm2d(64, track_running_stats=False),
                nn.ReLU(True)
            )
    
            self.encoder1 = nn.Sequential(
                nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1),
                nn.InstanceNorm2d(128, track_running_stats=False),
                nn.ReLU(True)
            )
    
            self.encoder2 = nn.Sequential(
                nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1),
                nn.InstanceNorm2d(256, track_running_stats=False),
                nn.ReLU(True)
            )
    
            blocks = []
            for _ in range(residual_blocks):
                block = ResnetBlock(256, 2)
                blocks.append(block)
    
            self.middle = nn.Sequential(*blocks)
    
            self.middle = nn.Sequential(
                # nn.Conv2d(256, 256, 3, 1, 1),
                # nn.InstanceNorm2d(256)
                nn.ReflectionPad2d(2),
                spectral_norm(nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=0, dilation=2, bias=not False), False),
                nn.InstanceNorm2d(256),
                nn.ReLU(True),
            )
    
            self.decoder = nn.Sequential(
                nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2, padding=1),
                nn.InstanceNorm2d(128, track_running_stats=False),
                nn.ReLU(True),
    
                nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1),
                nn.InstanceNorm2d(64, track_running_stats=False),
                nn.ReLU(True),
    
                nn.ReflectionPad2d(3),
                nn.Conv2d(in_channels=64, out_channels=3, kernel_size=7, padding=0),
            )
    
            self.kernel_pred = KernelConv(kernel_size=[3], sep_conv=False, core_bias=False)
    
            self.kpn_model = kpn_utils.create_generator()
    
            if init_weights:
                self.init_weights()
    
        def forward(self, x):
            inputs = x.clone()
    
            x = self.encoder0(x) # 64*256*256
            x = self.encoder1(x) # 128*128*128
    
            kernels, kernels_img = self.kpn_model(inputs, x)
    
            x = self.encoder2(x) # 256*64*64
            x = self.kernel_pred(x, kernels, white_level=1.0, rate=1)
    
            x = self.middle(x) # 256*64*64
    
            x = self.decoder(x) # 3*256*256
    
            x = self.kernel_pred(x, kernels_img, white_level=1.0, rate=1)
    
            x = (torch.tanh(x) + 1) / 2
    
            return x
    
        def save_feature(self, x, name):
            x = x.cpu().numpy()
            np.save('./result/{}'.format(name), x)
    
    
    class ResnetBlock(nn.Module):
        def __init__(self, dim, dilation=1, use_spectral_norm=False):
            super(ResnetBlock, self).__init__()
            self.conv_block = nn.Sequential(
                nn.ReflectionPad2d(dilation),
                spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation, bias=not use_spectral_norm), use_spectral_norm),
                nn.InstanceNorm2d(dim),
                nn.ReLU(True),
    
                nn.ReflectionPad2d(1),
                spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1, bias=not use_spectral_norm), use_spectral_norm),
                nn.InstanceNorm2d(dim),
            )
    
        def forward(self, x):
            out = x + self.conv_block(x)
    
            return out
    
    
    def spectral_norm(module, mode=True):
        if mode:
            return nn.utils.spectral_norm(module)
    
        return module
    
    
    
    if __name__ == '__main__':
        import pdb
        x = torch.randn(1, 4, 128, 128)
    
        model = InpaintGenerator()
        model.eval()
        torch.set_grad_enabled(False)
        torch.onnx.export(
            model, x, 'test.onnx',
            opset_version=14,  # 9 ~ 14
            input_names=['x'],
            output_names=['output'],
        )


**Code File 2: ./kpn/utils.py**

    import json
    import math
    import os
    
    import cv2
    import numpy as np
    import skimage
    import torch
    
    import kpn.network as network
    
    
    # ----------------------------------------
    #                 Network
    # ----------------------------------------
    def create_generator():
        generator = network.KPN()
        return generator
        
    def load_dict(process_net, pretrained_net):
        # Get the dict from pre-trained network
        pretrained_dict = pretrained_net
        # Get the dict from processing network
        process_dict = process_net.state_dict()
        # Delete the extra keys of pretrained_dict that do not belong to process_dict
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in process_dict}
        # Update process_dict using pretrained_dict
        process_dict.update(pretrained_dict)
        # Load the updated dict to processing network
        process_net.load_state_dict(process_dict)
        return process_net
    
    
    def save_model(config, iteration, generator, t=''):
        model_name = '{}_KPN_bs_{}_{}.pth'.format(iteration, config.BATCH_SIZE, t)
        save_model_path = os.path.join(config.kpn_model_save_path)
        if not os.path.exists(save_model_path):
            os.mkdir(save_model_path)
        save_model_path = os.path.join(save_model_path, model_name)
    
        if len(config.GPU) > 1:
            torch.save(generator.module.state_dict(), save_model_path)
            print('mul_gpu_The trained model is successfully saved at iteration {}'.format(iteration))
        else:
            torch.save(generator.state_dict(), save_model_path)
            print('The trained model is successfully saved at iteration {}'.format(iteration))
    
    
    # ----------------------------------------
    #    Validation and Sample at training
    # ----------------------------------------
    def save_sample_png(sample_folder, sample_name, img_list, name_list, pixel_max_cnt = 255, height = -1, width = -1):
        if not os.path.exists(sample_folder):
            os.mkdir(sample_folder)
    
        # Save image one-by-one
        for i in range(len(img_list)):
            img = img_list[i]
            # Recover normalization
            img = img * 255.0
    
            # Process img_copy and do not destroy the data of img
            #print(img.size())
            img_copy = img.clone().data.permute(0, 2, 3, 1).cpu().numpy()
            img_copy = np.clip(img_copy, 0, pixel_max_cnt)
            img_copy = img_copy.astype(np.uint8)[0, :, :, :]
            img_copy = cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB)
            if (height != -1) and (width != -1):
                img_copy = cv2.resize(img_copy, (width, height))
            # Save to certain path
            save_img_name = sample_name + '_' + name_list[i] + '.png'
            save_img_path = os.path.join(sample_folder, save_img_name)
    
            aa = img_copy[img_copy > 255]
            b = img_copy[img_copy < 0]
    
            cv2.imwrite(save_img_path, img_copy)
    
        return img_copy
    
    def save_sample_png_test(sample_folder, sample_name, img_list, name_list, pixel_max_cnt = 255):
        # Save image one-by-one
        for i in range(len(img_list)):
            img = img_list[i]
            # Recover normalization
            img = img * 255.0
            # Process img_copy and do not destroy the data of img
            img_copy = img.clone().data.permute(0, 2, 3, 1).cpu().numpy()
            img_copy = np.clip(img_copy, 0, pixel_max_cnt)
            img_copy = img_copy.astype(np.uint8)[0, :, :, :]
            img_copy = img_copy.astype(np.float32)
            img_copy = cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB)
            # Save to certain path
            save_img_name = sample_name + '_' + name_list[i] + '.png'
            save_img_path = os.path.join(sample_folder, save_img_name)
            cv2.imwrite(save_img_path, img_copy)
    
    def recover_process(img, height = -1, width = -1):
        img = img * 255.0
        img_copy = img.clone().data.permute(0, 2, 3, 1).cpu().numpy()
        img_copy = np.clip(img_copy, 0, 255)
        img_copy = img_copy.astype(np.uint8)[0, :, :, :]
        img_copy = img_copy.astype(np.float32)
        img_copy = cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB)
        if (height != -1) and (width != -1):
            img_copy = cv2.resize(img_copy, (width, height))
        return img_copy
    
    def psnr(pred, target):
        #print(pred.shape)
        #print(target.shape)
        mse = np.mean( (pred - target) ** 2 )
        if mse == 0:
            return 100
        PIXEL_MAX = 255.0
        return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
    
    
    def grey_psnr(pred, target, pixel_max_cnt = 255):
        pred = torch.sum(pred, dim = 0)
        target = torch.sum(target, dim = 0)
        mse = torch.mul(target - pred, target - pred)
        rmse_avg = (torch.mean(mse).item()) ** 0.5
        p = 20 * np.log10(pixel_max_cnt * 3 / rmse_avg)
        return p
    
    def ssim(pred, target):
        pred = pred.clone().data.permute(0, 2, 3, 1).cpu().numpy()
        target = target.clone().data.permute(0, 2, 3, 1).cpu().numpy()
        target = target[0]
        pred = pred[0]
        ssim = skimage.measure.compare_ssim(target, pred, multichannel = True)
        return ssim
    
    # ----------------------------------------
    #             PATH processing
    # ----------------------------------------
    def check_path(path):
        if not os.path.exists(path):
            os.makedirs(path)
    
    def savetxt(name, loss_log):
        np_loss_log = np.array(loss_log)
        np.savetxt(name, np_loss_log)
    
    
    #rain100H/L / SPA
    def get_files(path):
        if path is None:
            return []
        with open(path, 'r') as j:
            f_list = json.load(j)
            return f_list
    
    def get_jpgs(path):
        # read a folder, return the image name
        ret = [] 
        for root, dirs, files in os.walk(path):
            for filespath in files:
                ret.append(filespath)
        return ret
        
    def get_last_2paths(path):
        # read a folder, return the image name
        ret = [] 
        for root, dirs, files in os.walk(path):
            for filespath in files:
                if filespath[-4:] == '.png':
                    wholepath = os.path.join(root, filespath)
                    last_2paths = os.path.join(wholepath.split('/')[-2], wholepath.split('/')[-1])
                    ret.append(last_2paths)
        return ret
        
    def text_readlines(filename):
        # Try to read a txt file and return a list.Return [] if there was a mistake.
        try:
            file = open(filename, 'r')
        except IOError:
            error = []
            return error
        content = file.readlines()
        # This for loop deletes the EOF (like \n)
        for i in range(len(content)):
            content[i] = content[i][:len(content[i])-1]
        file.close()
        return content
    
    def text_save(content, filename, mode = 'a'):
        # save a list to a txt
        # Try to save a list variable in txt file.
        file = open(filename, mode)
        for i in range(len(content)):
            file.write(str(content[i]))
        file.close()


**Code File 3: ./kpn/network.py**

    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    import numpy as np
    
    # ----------------------------------------
    #         Initialize the networks
    # ----------------------------------------
    def weights_init(net, init_type = 'normal', init_gain = 0.02):
    
        def init_func(m):
            classname = m.__class__.__name__
            if hasattr(m, 'weight') and classname.find('Conv') != -1:
                if init_type == 'normal':
                    torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
                elif init_type == 'xavier':
                    torch.nn.init.xavier_normal_(m.weight.data, gain = init_gain)
                elif init_type == 'kaiming':
                    torch.nn.init.kaiming_normal_(m.weight.data, a = 0, mode = 'fan_in')
                elif init_type == 'orthogonal':
                    torch.nn.init.orthogonal_(m.weight.data, gain = init_gain)
                else:
                    raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            elif classname.find('BatchNorm2d') != -1:
                torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
                torch.nn.init.constant_(m.bias.data, 0.0)
    
        # apply the initialization function <init_func>
        print('initialize network with %s type' % init_type)
        net.apply(init_func)
    
    # ----------------------------------------
    #      Kernel Prediction Network (KPN)
    # ----------------------------------------
    class Basic(nn.Module):
        def __init__(self, in_ch, out_ch, g=16, channel_att=False, spatial_att=False):
            super(Basic, self).__init__()
            self.channel_att = channel_att
            self.spatial_att = spatial_att
            self.conv1 = nn.Sequential(
                    nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),
                    # nn.BatchNorm2d(out_ch),
                    nn.ReLU(),
                    nn.Conv2d(in_channels=out_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),
                    # nn.BatchNorm2d(out_ch),
                    nn.ReLU(),
                    nn.Conv2d(in_channels=out_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1),
                    # nn.BatchNorm2d(out_ch),
                    nn.ReLU()
                )
    
            if channel_att:
                self.att_c = nn.Sequential(
                    nn.Conv2d(2*out_ch, out_ch//g, 1, 1, 0),
                    nn.ReLU(),
                    nn.Conv2d(out_ch//g, out_ch, 1, 1, 0),
                    nn.Sigmoid()
                )
            if spatial_att:
                self.att_s = nn.Sequential(
                    nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3),
                    nn.Sigmoid()
                )
    
        def forward(self, data):
            """
            Forward function.
            :param data:
            :return: tensor
            """
            fm = self.conv1(data)
            if self.channel_att:
                # fm_pool = F.adaptive_avg_pool2d(fm, (1, 1)) + F.adaptive_max_pool2d(fm, (1, 1))
                fm_pool = torch.cat([F.adaptive_avg_pool2d(fm, (1, 1)), F.adaptive_max_pool2d(fm, (1, 1))], dim=1)
                att = self.att_c(fm_pool)
                fm = fm * att
            if self.spatial_att:
                fm_pool = torch.cat([torch.mean(fm, dim=1, keepdim=True), torch.max(fm, dim=1, keepdim=True)[0]], dim=1)
                att = self.att_s(fm_pool)
                fm = fm * att
            return fm
    
    class KPN(nn.Module):
        def __init__(self, kernel_size=[3], sep_conv=False, channel_att=False, spatial_att=False, upMode='bilinear', core_bias=False):
            super(KPN, self).__init__()
            self.upMode = upMode
            self.core_bias = core_bias
            self.kernel_size = kernel_size
    
            in_channel = 4
            out_channel = 64 * (self.kernel_size[0] ** 2)
    
            self.conv1 = Basic(in_channel, 64, channel_att=False, spatial_att=False) # 256*256
            self.conv2 = Basic(64, 128, channel_att=False, spatial_att=False)        # 128*128
            self.conv3 = Basic(128 + 128, 256, channel_att=False, spatial_att=False)  # 64*64
    
            self.conv4 = Basic(256, 512, channel_att=False, spatial_att=False)
    
            self.conv7 = Basic(256 + 512, 256, channel_att=channel_att, spatial_att=spatial_att)
            self.conv8 = Basic(256 + 256, 128, channel_att=channel_att, spatial_att=spatial_att)
            self.conv9 = Basic(128 + 64, 64, channel_att=channel_att, spatial_att=spatial_att)
    
            self.kernels = nn.Conv2d(256, out_channel, 1, 1, 0)
    
            out_channel_img = 3 * (self.kernel_size[0] ** 2)
            self.core_img = nn.Conv2d(64, out_channel_img, 1, 1, 0)
    
            self.kernel_pred = KernelConv(kernel_size, sep_conv, self.core_bias)
    
            self.conv_final = nn.Conv2d(in_channels=12, out_channels=3, kernel_size=3, stride=1, padding=1)
    
            self.iteration = 0
    
        def forward(self, data_with_est, x):
    
            conv1 = self.conv1(data_with_est) #64*256*256
            conv2 = self.conv2(F.avg_pool2d(conv1, kernel_size=2, stride=2)) # 128*128*128
    
            conv2 = torch.cat([conv2, x], dim=1)
    
            conv3 = self.conv3(F.avg_pool2d(conv2, kernel_size=2, stride=2)) # 256*64*64
            kernels = self.kernels(conv3)
            kernels = kernels.unsqueeze(dim=0)
    
    #       kernels = F.interpolate(input=kernels, size=(256*9, 64, 64), mode='nearest')
            kernels = F.interpolate(input=kernels, size=(256*9, data_with_est.shape[-1]//4, data_with_est.shape[-2]//4), mode='nearest')
    
            kernels = kernels.squeeze(dim=0)
    
    
            conv4 = self.conv4(conv3)
    
            conv7 = self.conv7(torch.cat([conv3, conv4], dim=1))
            conv8 = self.conv8(torch.cat([conv2, F.interpolate(conv7, scale_factor=2, mode=self.upMode)], dim=1))
            conv9 = self.conv9(torch.cat([conv1, F.interpolate(conv8, scale_factor=2, mode=self.upMode)], dim=1))
            core_img = self.core_img(conv9)
    
            return kernels, core_img
    
    
    class KernelConv(nn.Module):
        """
        the class of computing prediction
        """
        def __init__(self, kernel_size=[5], sep_conv=False, core_bias=False):
            super(KernelConv, self).__init__()
            self.kernel_size = sorted(kernel_size)
            self.sep_conv = sep_conv
            self.core_bias = core_bias
    
        def _sep_conv_core(self, core, batch_size, N, color, height, width):
            """
            convert the sep_conv core to conv2d core
            2p --> p^2
            :param core: shape: batch*(N*2*K)*height*width
            :return:
            """
            kernel_total = sum(self.kernel_size)
            core = core.view(batch_size, N, -1, color, height, width)
            if not self.core_bias:
                core_1, core_2 = torch.split(core, kernel_total, dim=2)
            else:
                core_1, core_2, core_3 = torch.split(core, kernel_total, dim=2)
            # output core
            core_out = {}
            cur = 0
            for K in self.kernel_size:
                t1 = core_1[:, :, cur:cur + K, ...].view(batch_size, N, K, 1, 3, height, width)
                t2 = core_2[:, :, cur:cur + K, ...].view(batch_size, N, 1, K, 3, height, width)
                core_out[K] = torch.einsum('ijklno,ijlmno->ijkmno', [t1, t2]).view(batch_size, N, K * K, color, height, width)
                cur += K
            # it is a dict
            return core_out, None if not self.core_bias else core_3.squeeze()
    
        def _convert_dict(self, core, batch_size, N, color, height, width):
            """
            make sure the core to be a dict, generally, only one kind of kernel size is suitable for the func.
            :param core: shape: batch_size*(N*K*K)*height*width
            :return: core_out, a dict
            """
            core_out = {}
            core = core.view(batch_size, N, -1, color, height, width)
            core_out[self.kernel_size[0]] = core[:, :, 0:self.kernel_size[0]**2, ...]
            bias = None if not self.core_bias else core[:, :, -1, ...]
            return core_out, bias
    
        def forward(self, frames, core, white_level=1.0, rate=1):
            """
            compute the pred image according to core and frames
            :param frames: [batch_size, N, 3, height, width]
            :param core: [batch_size, N, dict(kernel), 3, height, width]
            :return:
            """
            if len(frames.size()) == 5:
                batch_size, N, color, height, width = frames.size()
            else:
                batch_size, N, height, width = frames.size()
                color = 1
                frames = frames.view(batch_size, N, color, height, width)
            if self.sep_conv:
                core, bias = self._sep_conv_core(core, batch_size, N, color, height, width)
            else:
                core, bias = self._convert_dict(core, batch_size, N, color, height, width)
            img_stack = []
            pred_img = []
            kernel = self.kernel_size[::-1]
            for index, K in enumerate(kernel):
                if not img_stack:
                    padding_num = (K//2) * rate
                    frame_pad = F.pad(frames, [padding_num, padding_num, padding_num, padding_num])
                    for i in range(0, K):
                        for j in range(0, K):
                            img_stack.append(frame_pad[..., i*rate:i*rate + height, j*rate:j*rate + width])
                    img_stack = torch.stack(img_stack, dim=2)
                else:
                    k_diff = (kernel[index - 1] - kernel[index]) // 2
                    img_stack = img_stack[:, :, k_diff:-k_diff, ...]
                # print('img_stack:', img_stack.size())
                pred_img.append(torch.sum(
                    core[K].mul(img_stack), dim=2, keepdim=False
                ))
            pred_img = torch.stack(pred_img, dim=0)
            # print('pred_stack:', pred_img.size())
            pred_img_i = torch.mean(pred_img, dim=0, keepdim=False)
            #print("pred_img_i", pred_img_i.size())
            # N = 1
            pred_img_i = pred_img_i.squeeze(2)
            #print("pred_img_i", pred_img_i.size())
            # if bias is permitted
            if self.core_bias:
                if bias is None:
                    raise ValueError('The bias should not be None.')
                pred_img_i += bias
            # print('white_level', white_level.size())
            pred_img_i = pred_img_i / white_level
            #pred_img = torch.mean(pred_img_i, dim=1, keepdim=True)
            # print('pred_img:', pred_img.size())
            # print('pred_img_i:', pred_img_i.size())
            return pred_img_i
    
    class LossFunc(nn.Module):
        """
        loss function of KPN
        """
        def __init__(self, coeff_basic=1.0, coeff_anneal=1.0, gradient_L1=True, alpha=0.9998, beta=100):
            super(LossFunc, self).__init__()
            self.coeff_basic = coeff_basic
            self.coeff_anneal = coeff_anneal
            self.loss_basic = LossBasic(gradient_L1)
            self.loss_anneal = LossAnneal(alpha, beta)
    
        def forward(self, pred_img_i, pred_img, ground_truth, global_step):
            """
            forward function of loss_func
            :param frames: frame_1 ~ frame_N, shape: [batch, N, 3, height, width]
            :param core: a dict coverted by ......
            :param ground_truth: shape [batch, 3, height, width]
            :param global_step: int
            :return: loss
            """
            return self.coeff_basic * self.loss_basic(pred_img, ground_truth), self.coeff_anneal * self.loss_anneal(global_step, pred_img_i, ground_truth)
    
    class LossBasic(nn.Module):
        """
        Basic loss function.
        """
        def __init__(self, gradient_L1=True):
            super(LossBasic, self).__init__()
            self.l1_loss = nn.L1Loss()
            self.l2_loss = nn.MSELoss()
            self.gradient = TensorGradient(gradient_L1)
    
        def forward(self, pred, ground_truth):
            return self.l2_loss(pred, ground_truth) + \
                   self.l1_loss(self.gradient(pred), self.gradient(ground_truth))
    
    class LossAnneal(nn.Module):
        """
        anneal loss function
        """
        def __init__(self, alpha=0.9998, beta=100):
            super(LossAnneal, self).__init__()
            self.global_step = 0
            self.loss_func = LossBasic(gradient_L1=True)
            self.alpha = alpha
            self.beta = beta
    
        def forward(self, global_step, pred_i, ground_truth):
            """
            :param global_step: int
            :param pred_i: [batch_size, N, 3, height, width]
            :param ground_truth: [batch_size, 3, height, width]
            :return:
            """
            loss = 0
            for i in range(pred_i.size(1)):
                loss += self.loss_func(pred_i[:, i, ...], ground_truth)
            loss /= pred_i.size(1)
            return self.beta * self.alpha ** global_step * loss
    
    class TensorGradient(nn.Module):
        """
        the gradient of tensor
        """
        def __init__(self, L1=True):
            super(TensorGradient, self).__init__()
            self.L1 = L1
    
        def forward(self, img):
            w, h = img.size(-2), img.size(-1)
            l = F.pad(img, [1, 0, 0, 0])
            r = F.pad(img, [0, 1, 0, 0])
            u = F.pad(img, [0, 0, 1, 0])
            d = F.pad(img, [0, 0, 0, 1])
            if self.L1:
                return torch.abs((l - r)[..., 0:w, 0:h]) + torch.abs((u - d)[..., 0:w, 0:h])
            else:
                return torch.sqrt(
                    torch.pow((l - r)[..., 0:w, 0:h], 2) + torch.pow((u - d)[..., 0:w, 0:h], 2)
                )
    
    if __name__ == '__main__':
        
        kpn = KPN().cuda()
        a = torch.randn(4, 3, 224, 224).cuda()
        b = kpn(a, a)
        print(b.shape)



### Versions

pip install opencv==4.7.0.72
pip install numpy==1.24.2
pip install  skimage==0.20.0

torch              1.13.0+cu117
torchvision        0.14.0+cu117