import sys
_module = sys.modules[__name__]
del sys
demo = _module
eval = _module
my_eval = _module
my_test = _module
config = _module
net = _module
tracker = _module
utils = _module
toolkit = _module
datasets = _module
dataset = _module
got10k = _module
lasot = _module
nfs = _module
otb = _module
trackingnet = _module
uav = _module
video = _module
vot = _module
evaluation = _module
ar_benchmark = _module
eao_benchmark = _module
f1_benchmark = _module
ope_benchmark = _module
misc = _module
statistics = _module
visualization = _module
draw_eao = _module
draw_f1 = _module
draw_success_precision = _module
draw_utils = _module
demo = _module
hp_search = _module
test = _module
train = _module
cal_macs_params = _module
cal_speed = _module
dtb70 = _module
tcolor128 = _module
uav123 = _module
uavdt = _module
vid = _module
visdrone = _module
experiments = _module
trackers = _module
identity_tracker = _module
ioutils = _module
metrics = _module
viz = _module
nanotrack = _module
core = _module
xcorr = _module
augmentation = _module
dataset = _module
point_target = _module
models = _module
backbone = _module
alexnet = _module
mobile_v3 = _module
head = _module
ban_v1 = _module
ban_v2 = _module
init_weight = _module
iou_loss = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
nano_tracker = _module
tracker_builder = _module
average_meter = _module
bbox = _module
distributed = _module
log_helper = _module
lr_scheduler = _module
model_load = _module
point = _module
pytorch2onnx = _module
setup = _module
logger = _module
demo = _module
hp_search = _module
my_test = _module
my_train = _module
test = _module
test_epochs = _module
tune = _module
siamban = _module
xcorr = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
ban = _module
init_weight = _module
iou_loss = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siamban_tracker = _module
distributed = _module
lr_scheduler = _module
model_load = _module
logger = _module
demo = _module
dist_train = _module
hp_search = _module
my_test = _module
my_train = _module
test = _module
train = _module
siamcar = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
car_head = _module
init_weight = _module
loss_car = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siamcar_tracker = _module
distributed = _module
location_grid = _module
lr_scheduler = _module
misc = _module
model_load = _module
xcorr = _module
speed = _module
logger = _module
siamfc = _module
backbones = _module
bbox_util = _module
config = _module
datasets = _module
heads = _module
losses = _module
modules = _module
network = _module
tracker = _module
transforms = _module
utils = _module
my_train = _module
siamrpn = _module
backbones = _module
dataset = _module
heads = _module
loss = _module
modules = _module
network = _module
tracker = _module
transforms = _module
utils = _module
visual = _module
backbones = _module
config = _module
datasets = _module
heads = _module
losses = _module
modules = _module
network = _module
tracker = _module
transforms = _module
utils = _module
demo = _module
dist_train = _module
hp_search = _module
hp_search2 = _module
my_test = _module
my_test2 = _module
my_train = _module
test = _module
train = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
car_head = _module
init_weight = _module
loss_car = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siamcar_tracker = _module
distributed = _module
location_grid = _module
lr_scheduler = _module
misc = _module
model_load = _module
xcorr = _module
speed = _module
logger = _module
visualize_siamfcpp_training_data = _module
static_image = _module
sot_video = _module
resources = _module
static_img_example = _module
get_image = _module
dist_train = _module
contrib_module = _module
builder = _module
contrib_module_base = _module
contrib_module_impl = _module
inherited_contrib_module_impl = _module
template_module = _module
template_module_base = _module
template_module_impl = _module
inherited_template_module_impl = _module
hpo = _module
my_test = _module
my_train = _module
test = _module
train = _module
siamfcpp = _module
data = _module
adaptor_dataset = _module
builder = _module
datapipeline = _module
datapipeline_base = _module
datapipeline_impl = _module
regular_datapipeline = _module
dataset_base = _module
dataset_impl = _module
coco = _module
davis = _module
ilsvrc_det = _module
ilsvrc_vid = _module
ytb_vos = _module
filter = _module
filter_base = _module
filter_impl = _module
track_pair_filter = _module
sampler = _module
sampler_base = _module
sampler_impl = _module
track_pair_sampler = _module
target = _module
target_base = _module
target_impl = _module
densebox_target = _module
identity_target = _module
sat_mask_target = _module
make_densebox_target = _module
transformer = _module
transformer_base = _module
transformer_impl = _module
common_transformer = _module
random_crop_transformer = _module
crop_track_pair = _module
filter_box = _module
engine = _module
monitor = _module
monitor_base = _module
monitor_impl = _module
monitor_template = _module
seg_metric = _module
tensorboard_logger = _module
track_info = _module
tester = _module
tester_base = _module
tester_impl = _module
davis = _module
dtb70 = _module
got10k = _module
lasot = _module
otb = _module
trackingnet = _module
uav123 = _module
uavdt = _module
got_benchmark_helper = _module
visdrone = _module
vot = _module
trainer = _module
trainer_base = _module
trainer_impl = _module
distributed_regular_trainer = _module
distributed_sat_trainer = _module
regular_trainer = _module
davis_benchmark = _module
benckmark_helpler = _module
davis_evaluation = _module
results = _module
evaluation_method = _module
got_benchmark = _module
vot_benchmark = _module
bbox_helper = _module
benchmark_helper = _module
pysot = _module
model = _module
backbone_base = _module
backbone_impl = _module
alexnet_bn = _module
googlenet = _module
resnet = _module
shufflenet_v2 = _module
tinyconv = _module
builder = _module
builder = _module
common_block = _module
loss_base = _module
loss_impl = _module
iou_loss = _module
multi_bceloss = _module
sigmoid_ce_centerness = _module
sigmoid_ce_retina = _module
utils = _module
module_base = _module
task_head = _module
builder = _module
taskhead_base = _module
taskhead_impl = _module
track_head = _module
vos_head = _module
task_model = _module
builder = _module
taskmodel_base = _module
taskmodel_impl = _module
sat_vos = _module
siamese_track = _module
load_state = _module
optim = _module
builder = _module
grad_modifier = _module
grad_modifier_base = _module
grad_modifier_impl = _module
dynamic_freezer = _module
freeze = _module
optimizer = _module
builder = _module
optimizer_base = _module
optimizer_impl = _module
sgd = _module
lr_multiply = _module
lr_policy = _module
pipeline = _module
builder = _module
pipeline_base = _module
segmenter_impl = _module
sat_pipeline = _module
tracker_impl = _module
siamfcpp_osdet = _module
siamfcpp_track = _module
siamfcpp_track_multi_temp = _module
crop = _module
misc = _module
dist_utils = _module
image = _module
path = _module
torch_module = _module
SiameseNetworkTest = _module
SiameseNetworkTrain = _module
SiameseUtil = _module
demo = _module
dist_train = _module
hp_search = _module
my_demo = _module
my_test = _module
my_train = _module
my_train_base = _module
my_train_sharp = _module
test = _module
siamrpnpp = _module
xcorr = _module
anchor_target = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
mask = _module
rpn = _module
init_weight = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siammask_tracker = _module
siamrpn_tracker = _module
siamrpnlt_tracker = _module
anchor = _module
distributed = _module
lr_scheduler = _module
model_load = _module
logger = _module
demo = _module
dist_train = _module
hp_search = _module
my_test = _module
my_train = _module
test = _module
xcorr = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
mask = _module
rpn = _module
init_weight = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siamrpn_tracker = _module
distributed = _module
lr_scheduler = _module
model_load = _module
logger = _module
dist_train = _module
my_test = _module
my_train = _module
dataset = _module
distributed = _module
loss = _module
network = _module
tracker = _module
transforms = _module
utils = _module
demo = _module
dist_train = _module
hp_search = _module
my_test = _module
my_train = _module
test = _module
xcorr = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
mask = _module
rpn = _module
init_weight = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
siamrpn_tracker = _module
distributed = _module
lr_scheduler = _module
model_load = _module
logger = _module
demo = _module
hp_search = _module
my_test = _module
my_train = _module
test = _module
tune = _module
logger = _module
trtr = _module
xcorr = _module
dataset = _module
alexnet = _module
mobile_v2 = _module
resnet_atrous = _module
ban = _module
trans_head = _module
init_weight = _module
iou_loss = _module
loss = _module
model_builder = _module
neck = _module
neck = _module
base_tracker = _module
distributed = _module
lr_scheduler = _module
model_load = _module
demo = _module
my_test = _module
net = _module
tracker = _module
utils = _module
basenet = _module
config_upd = _module
create_template = _module
net_upd = _module
tracker_upd = _module
train_upd = _module
utils_upd = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import torch


import numpy as np


import torch.nn as nn


import torch.nn.functional as F


from torch.autograd import Variable


import logging


import time


import math


import random


from torch.utils.data import DataLoader


from torch.utils.tensorboard import SummaryWriter


from torch.nn.utils import clip_grad_norm_


from torch.utils.data.distributed import DistributedSampler


from torch.utils.data import Dataset


from torch import nn


import torch.distributed as dist


from torch.optim.lr_scheduler import _LRScheduler


from collections import namedtuple


from collections import OrderedDict


import torch.optim as optim


from torch.optim.lr_scheduler import ExponentialLR


from torch.nn.parallel import DistributedDataParallel


import numbers


import torchvision.transforms as transforms


import torchvision


import pandas as pd


from torch.optim.lr_scheduler import StepLR


from sklearn.model_selection import train_test_split


from torchvision import datasets


from torchvision import transforms


from torchvision import utils


from torchvision.models import alexnet


import torch.nn


import functools


from torch.multiprocessing import Pool


from torch.multiprocessing import Manager


import torch.multiprocessing as mp


from itertools import chain


import torch.multiprocessing


from typing import Dict


from typing import List


from torch.utils.data.dataloader import default_collate


from abc import ABCMeta


from copy import deepcopy


import itertools


from collections import Mapping


from torch.utils.tensorboard.writer import SummaryWriter


import copy


from typing import Tuple


from collections import defaultdict


from typing import Any


from typing import Iterable


from torch import optim


import re


from torch.optim.optimizer import Optimizer


import collections.abc


import torchvision.datasets as dataset


import matplotlib.pyplot as plt


import torchvision.utils


from torch.autograd import Function


import torch.nn.init as init


from typing import Optional


from torch import Tensor


from torch.utils.data import dataloader


import torch.backends.cudnn as cudnn


class SiamRPN(nn.Module):

    def __init__(self, size=2, feature_out=512, anchor=5):
        configs = [3, 96, 256, 384, 384, 256]
        configs = list(map(lambda x: 3 if x == 3 else x * size, configs))
        feat_in = configs[-1]
        super(SiamRPN, self).__init__()
        self.featureExtract = nn.Sequential(nn.Conv2d(configs[0], configs[1], kernel_size=11, stride=2), nn.BatchNorm2d(configs[1]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(configs[1], configs[2], kernel_size=5), nn.BatchNorm2d(configs[2]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(configs[2], configs[3], kernel_size=3), nn.BatchNorm2d(configs[3]), nn.ReLU(inplace=True), nn.Conv2d(configs[3], configs[4], kernel_size=3), nn.BatchNorm2d(configs[4]), nn.ReLU(inplace=True), nn.Conv2d(configs[4], configs[5], kernel_size=3), nn.BatchNorm2d(configs[5]))
        self.anchor = anchor
        self.feature_out = feature_out
        self.conv_r1 = nn.Conv2d(feat_in, feature_out * 4 * anchor, 3)
        self.conv_r2 = nn.Conv2d(feat_in, feature_out, 3)
        self.conv_cls1 = nn.Conv2d(feat_in, feature_out * 2 * anchor, 3)
        self.conv_cls2 = nn.Conv2d(feat_in, feature_out, 3)
        self.regress_adjust = nn.Conv2d(4 * anchor, 4 * anchor, 1)
        self.r1_kernel = []
        self.cls1_kernel = []
        self.cfg = {}

    def forward(self, x):
        x_f = self.featureExtract(x)
        return self.regress_adjust(F.conv2d(self.conv_r2(x_f), self.r1_kernel)), F.conv2d(self.conv_cls2(x_f), self.cls1_kernel)

    def featextract(self, x):
        x_f = self.featureExtract(x)
        return x_f

    def kernel(self, z_f):
        r1_kernel_raw = self.conv_r1(z_f)
        cls1_kernel_raw = self.conv_cls1(z_f)
        kernel_size = r1_kernel_raw.data.size()[-1]
        self.r1_kernel = r1_kernel_raw.view(self.anchor * 4, self.feature_out, kernel_size, kernel_size)
        self.cls1_kernel = cls1_kernel_raw.view(self.anchor * 2, self.feature_out, kernel_size, kernel_size)

    def temple(self, z):
        z_f = self.featureExtract(z)
        r1_kernel_raw = self.conv_r1(z_f)
        cls1_kernel_raw = self.conv_cls1(z_f)
        kernel_size = r1_kernel_raw.data.size()[-1]
        self.r1_kernel = r1_kernel_raw.view(self.anchor * 4, self.feature_out, kernel_size, kernel_size)
        self.cls1_kernel = cls1_kernel_raw.view(self.anchor * 2, self.feature_out, kernel_size, kernel_size)


class SiamRPNBIG(SiamRPN):

    def __init__(self):
        super(SiamRPNBIG, self).__init__(size=2)
        self.cfg = {'lr': 0.295, 'window_influence': 0.42, 'penalty_k': 0.055, 'instance_size': 271, 'adaptive': True}


class SiamRPNVOT(SiamRPN):

    def __init__(self):
        super(SiamRPNVOT, self).__init__(size=1, feature_out=256)
        self.cfg = {'lr': 0.45, 'window_influence': 0.44, 'penalty_k': 0.04, 'instance_size': 271, 'adaptive': False}


class SiamRPNOTB(SiamRPN):

    def __init__(self):
        super(SiamRPNOTB, self).__init__(size=1, feature_out=256)
        self.cfg = {'lr': 0.3, 'window_influence': 0.4, 'penalty_k': 0.22, 'instance_size': 271, 'adaptive': False}


class AlexNet(nn.Module):
    configs = [3, 96, 256, 384, 384, 256]

    def __init__(self, width_mult=1):
        configs = list(map(lambda x: 3 if x == 3 else int(x * width_mult), AlexNet.configs))
        super(AlexNet, self).__init__()
        self.layer1 = nn.Sequential(nn.Conv2d(configs[0], configs[1], kernel_size=11, stride=2), nn.BatchNorm2d(configs[1]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True))
        self.layer2 = nn.Sequential(nn.Conv2d(configs[1], configs[2], kernel_size=5), nn.BatchNorm2d(configs[2]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True))
        self.layer3 = nn.Sequential(nn.Conv2d(configs[2], configs[3], kernel_size=3), nn.BatchNorm2d(configs[3]), nn.ReLU(inplace=True))
        self.layer4 = nn.Sequential(nn.Conv2d(configs[3], configs[4], kernel_size=3), nn.BatchNorm2d(configs[4]), nn.ReLU(inplace=True))
        self.layer5 = nn.Sequential(nn.Conv2d(configs[4], configs[5], kernel_size=3), nn.BatchNorm2d(configs[5]))
        self.feature_size = configs[5]

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        return x


class AlexNetLegacy(nn.Module):
    configs = [3, 96, 256, 384, 384, 256]

    def __init__(self, width_mult=1):
        configs = list(map(lambda x: 3 if x == 3 else int(x * width_mult), AlexNet.configs))
        super(AlexNetLegacy, self).__init__()
        self.features = nn.Sequential(nn.Conv2d(configs[0], configs[1], kernel_size=11, stride=2), nn.BatchNorm2d(configs[1]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(configs[1], configs[2], kernel_size=5), nn.BatchNorm2d(configs[2]), nn.MaxPool2d(kernel_size=3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(configs[2], configs[3], kernel_size=3), nn.BatchNorm2d(configs[3]), nn.ReLU(inplace=True), nn.Conv2d(configs[3], configs[4], kernel_size=3), nn.BatchNorm2d(configs[4]), nn.ReLU(inplace=True), nn.Conv2d(configs[4], configs[5], kernel_size=3), nn.BatchNorm2d(configs[5]))
        self.feature_size = configs[5]

    def forward(self, x):
        x = self.features(x)
        return x


class h_sigmoid(nn.Module):

    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.hard_sigmoid = nn.Hardsigmoid(inplace=inplace)

    def forward(self, x):
        return self.hard_sigmoid(x)


class h_swish(nn.Module):

    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.hard_swish = nn.Hardswish(inplace=True)

    def forward(self, x):
        return self.hard_swish(x)


def _make_divisible(v, divisor, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


class SELayer(nn.Module):

    def __init__(self, channel, reduction=4):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(nn.Linear(channel, _make_divisible(channel // reduction, 8)), nn.ReLU(inplace=True), nn.Linear(_make_divisible(channel // reduction, 8), channel), h_sigmoid())

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y


class InvertedResidual(nn.Module):

    def __init__(self, inp, oup, stride, expand_ratio, dilation=1):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        self.use_res_connect = self.stride == 1 and inp == oup
        padding = 2 - stride
        if dilation > 1:
            padding = dilation
        self.conv = nn.Sequential(nn.Conv2d(inp, inp * expand_ratio, 1, 1, 0, bias=False), nn.BatchNorm2d(inp * expand_ratio), nn.ReLU6(inplace=True), nn.Conv2d(inp * expand_ratio, inp * expand_ratio, 3, stride, padding, dilation=dilation, groups=inp * expand_ratio, bias=False), nn.BatchNorm2d(inp * expand_ratio), nn.ReLU6(inplace=True), nn.Conv2d(inp * expand_ratio, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup))

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


def conv_3x3_bn(inp, oup, stride):
    return nn.Sequential(nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup), h_swish())


class MobileNetV3(nn.Module):

    def __init__(self, cfgs, mode, num_classes=1000, width_mult=1.0, used_layers=[2, 3, 4]):
        super(MobileNetV3, self).__init__()
        self.cfgs = cfgs
        assert mode in ['large', 'small']
        input_channel = _make_divisible(16 * width_mult, 8)
        layers = [conv_3x3_bn(3, input_channel, 2)]
        block = InvertedResidual
        for k, t, c, use_se, use_hs, s in self.cfgs:
            output_channel = _make_divisible(c * width_mult, 8)
            exp_size = _make_divisible(input_channel * t, 8)
            layers.append(block(input_channel, exp_size, output_channel, k, s, use_se, use_hs))
            input_channel = output_channel
        self.features = nn.Sequential(*layers)
        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()


class BAN(nn.Module):

    def __init__(self):
        super(BAN, self).__init__()

    def forward(self, z_f, x_f):
        raise NotImplementedError


def xcorr_fast(x, kernel):
    """group conv2d to calculate cross correlation, fast version
    """
    batch = kernel.size()[0]
    pk = kernel.view(-1, x.size()[1], kernel.size()[2], kernel.size()[3])
    px = x.view(1, -1, x.size()[2], x.size()[3])
    po = F.conv2d(px, pk, groups=batch)
    po = po.view(batch, -1, po.size()[2], po.size()[3])
    return po


class UPChannelBAN(BAN):

    def __init__(self, feature_in=256, cls_out_channels=2):
        super(UPChannelBAN, self).__init__()
        cls_output = cls_out_channels
        loc_output = 4
        self.template_cls_conv = nn.Conv2d(feature_in, feature_in * cls_output, kernel_size=3)
        self.template_loc_conv = nn.Conv2d(feature_in, feature_in * loc_output, kernel_size=3)
        self.search_cls_conv = nn.Conv2d(feature_in, feature_in, kernel_size=3)
        self.search_loc_conv = nn.Conv2d(feature_in, feature_in, kernel_size=3)
        self.loc_adjust = nn.Conv2d(loc_output, loc_output, kernel_size=1)

    def forward(self, z_f, x_f):
        cls_kernel = self.template_cls_conv(z_f)
        loc_kernel = self.template_loc_conv(z_f)
        cls_feature = self.search_cls_conv(x_f)
        loc_feature = self.search_loc_conv(x_f)
        cls = xcorr_fast(cls_feature, cls_kernel)
        loc = self.loc_adjust(xcorr_fast(loc_feature, loc_kernel))
        return cls, loc


def xcorr_depthwise(x, kernel):
    """depthwise cross correlation
    """
    batch = kernel.size(0)
    channel = kernel.size(1)
    x = x.view(1, batch * channel, x.size(2), x.size(3))
    kernel = kernel.view(batch * channel, 1, kernel.size(2), kernel.size(3))
    out = F.conv2d(x, kernel, groups=batch * channel)
    out = out.view(batch, channel, out.size(2), out.size(3))
    return out


class DepthwiseXCorr(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(DepthwiseXCorr, self).__init__()
        self.conv_kernel = nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, bias=False), nn.BatchNorm2d(in_channels), nn.ReLU(inplace=True))
        self.conv_search = nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, bias=False), nn.BatchNorm2d(in_channels), nn.ReLU(inplace=True))
        for modules in [self.conv_kernel, self.conv_search]:
            for m in modules.modules():
                if isinstance(m, nn.Conv2d):
                    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                    m.weight.data.normal_(0, math.sqrt(2.0 / n))
                    if m.bias is not None:
                        m.bias.data.zero_()
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()
                elif isinstance(m, nn.Linear):
                    m.weight.data.normal_(0, 0.01)
                    m.bias.data.zero_()

    def forward(self, kernel, search):
        kernel = self.conv_kernel(kernel)
        search = self.conv_search(search)
        feature = xcorr_depthwise(search, kernel)
        return feature


class CAModule(nn.Module):
    """Channel attention module"""

    def __init__(self, channels=64, reduction=1):
        super(CAModule, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, padding=0)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1, padding=0)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        module_input = x
        x = self.avg_pool(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return module_input * x


def xcorr_pixelwise(x, kernel):
    """Pixel-wise correlation (implementation by matrix multiplication)
    The speed is faster because the computation is vectorized"""
    b, c, h, w = x.size()
    kernel_mat = kernel.view((b, c, -1)).transpose(1, 2)
    x_mat = x.view((b, c, -1))
    return torch.matmul(kernel_mat, x_mat).view((b, -1, h, w))


class PixelwiseXCorr(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(PixelwiseXCorr, self).__init__()
        self.CA_layer = CAModule(channels=64)
        self.conv_kernel = nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1), nn.BatchNorm2d(in_channels))
        self.conv_search = nn.Sequential(nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1), nn.BatchNorm2d(in_channels))
        for modules in [self.conv_kernel, self.conv_search]:
            for m in modules.modules():
                if isinstance(m, nn.Conv2d):
                    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                    m.weight.data.normal_(0, math.sqrt(2.0 / n))
                    if m.bias is not None:
                        m.bias.data.zero_()
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()
                elif isinstance(m, nn.Linear):
                    m.weight.data.normal_(0, 0.01)
                    m.bias.data.zero_()

    def forward(self, kernel, search):
        kernel = self.conv_kernel(kernel)
        search = self.conv_search(search)
        feature = xcorr_pixelwise(search, kernel)
        corr = self.CA_layer(feature)
        return corr


class NestedTensor(object):

    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        cast_tensor = self.tensors
        mask = self.mask
        if mask is not None:
            assert mask is not None
            cast_mask = mask
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):
        return str(self.tensors)


class PositionEmbeddingSine(nn.Module):
    """
    This is a more standard version of the position embedding, very similar to the one
    used by the Attention is all you need paper, generalized to work on images.
    """

    def __init__(self, num_pos_feats=128, temperature=10000, normalize=True, scale=None):
        super().__init__()
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError('normalize should be True if scale is passed')
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale
        self.segment_embdded_factor = 0.0

    def forward(self, tensor_list: NestedTensor, multi_frame=False):
        x = tensor_list
        b, c, h, w = tensor_list.shape
        not_mask = torch.ones([b, h, w], device=tensor_list.device)
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            eps = 1e-06
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        if multi_frame:
            assert len(pos) == 2
            pos[-1].add_(self.segment_embdded_factor)
        return pos


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


class TransformerDecoder(nn.Module):

    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate

    def forward(self, tgt, memory, tgt_mask: Optional[Tensor]=None, memory_mask: Optional[Tensor]=None, tgt_key_padding_mask: Optional[Tensor]=None, memory_key_padding_mask: Optional[Tensor]=None, encoder_pos: Optional[Tensor]=None, decoder_pos: Optional[Tensor]=None):
        output = tgt
        intermediate = []
        for layer in self.layers:
            output = layer(output, memory, tgt_mask=tgt_mask, memory_mask=memory_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask, encoder_pos=encoder_pos, decoder_pos=decoder_pos)
            if self.return_intermediate:
                intermediate.append(self.norm(output))
        if self.norm is not None:
            output = self.norm(output)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)
        if self.return_intermediate:
            return torch.stack(intermediate)
        return output.unsqueeze(0)


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == 'relu':
        return F.relu
    if activation == 'gelu':
        return F.gelu
    if activation == 'glu':
        return F.glu
    raise RuntimeError(f'activation should be relu/gelu, not {activation}.')


class TransformerDecoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory, tgt_mask: Optional[Tensor]=None, memory_mask: Optional[Tensor]=None, tgt_key_padding_mask: Optional[Tensor]=None, memory_key_padding_mask: Optional[Tensor]=None, encoder_pos: Optional[Tensor]=None, decoder_pos: Optional[Tensor]=None):
        q = k = self.with_pos_embed(tgt, decoder_pos)
        tgt2, attn_weight_map = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2, attn_weight_map = self.multihead_attn(query=self.with_pos_embed(tgt, decoder_pos), key=self.with_pos_embed(memory, encoder_pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

    def forward_pre(self, tgt, memory, tgt_mask: Optional[Tensor]=None, memory_mask: Optional[Tensor]=None, tgt_key_padding_mask: Optional[Tensor]=None, memory_key_padding_mask: Optional[Tensor]=None, encoder_pos: Optional[Tensor]=None, decoder_pos: Optional[Tensor]=None):
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, decoder_pos)
        tgt2, attn_weight_map = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        tgt2, attn_weight_map = self.multihead_attn(query=self.with_pos_embed(tgt2, decoder_pos), key=self.with_pos_embed(memory, encoder_pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt

    def forward(self, tgt, memory, tgt_mask: Optional[Tensor]=None, memory_mask: Optional[Tensor]=None, tgt_key_padding_mask: Optional[Tensor]=None, memory_key_padding_mask: Optional[Tensor]=None, encoder_pos: Optional[Tensor]=None, decoder_pos: Optional[Tensor]=None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, encoder_pos, decoder_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, encoder_pos, decoder_pos)


class TransformerEncoder(nn.Module):

    def __init__(self, encoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src, mask: Optional[Tensor]=None, src_key_padding_mask: Optional[Tensor]=None, pos: Optional[Tensor]=None):
        output = src
        for layer in self.layers:
            output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos)
        if self.norm is not None:
            output = self.norm(output)
        return output


class TransformerEncoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.0, activation='relu', normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, src, src_mask: Optional[Tensor]=None, src_key_padding_mask: Optional[Tensor]=None, pos: Optional[Tensor]=None):
        q = k = self.with_pos_embed(src, pos)
        src2, attn_weight_map = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

    def forward_pre(self, src, src_mask: Optional[Tensor]=None, src_key_padding_mask: Optional[Tensor]=None, pos: Optional[Tensor]=None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        return src

    def forward(self, src, src_mask: Optional[Tensor]=None, src_key_padding_mask: Optional[Tensor]=None, pos: Optional[Tensor]=None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)


class Transformer(nn.Module):

    def __init__(self, cfg, d_model=256, nhead=8, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=512, dropout=0.0, activation='relu', normalize_before=False, return_intermediate_dec=False):
        super().__init__()
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, return_intermediate=return_intermediate_dec)
        self.pos_encoder = PositionEmbeddingSine()
        self._reset_parameters()
        self.d_model = d_model
        self.nhead = nhead
        self.memory = []

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, template_src, search_src):
        template_mask, search_mask = None, None
        template_pos_embed = self.pos_encoder(template_src)
        search_pos_embed = self.pos_encoder(search_src)
        template_src = template_src.flatten(2).permute(2, 0, 1)
        template_pos_embed = template_pos_embed.flatten(2).permute(2, 0, 1)
        if template_mask is not None:
            template_mask = template_mask.flatten(1)
        memory = self.encoder(template_src, src_key_padding_mask=template_mask, pos=template_pos_embed)
        search_src = search_src.flatten(2).permute(2, 0, 1)
        search_pos_embed = search_pos_embed.flatten(2).permute(2, 0, 1)
        if template_mask is not None:
            search_mask = search_mask.flatten(1)
        hs = self.decoder(search_src, memory, memory_key_padding_mask=template_mask, tgt_key_padding_mask=search_mask, encoder_pos=template_pos_embed, decoder_pos=search_pos_embed)
        return hs.transpose(1, 2)


class DepthwiseBAN(BAN):

    def __init__(self, in_channels=96, out_channels=96, weighted=False):
        super(DepthwiseBAN, self).__init__()
        self.cls_dw = DepthwiseXCorr(in_channels, out_channels)
        self.reg_dw = DepthwiseXCorr(in_channels, out_channels)
        cls_tower = []
        bbox_tower = []
        for i in range(cfg.TRAIN.NUM_CONVS):
            cls_tower.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1))
            cls_tower.append(nn.GroupNorm(32, in_channels))
            cls_tower.append(nn.ReLU())
            bbox_tower.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1))
            bbox_tower.append(nn.GroupNorm(32, in_channels))
            bbox_tower.append(nn.ReLU())
        self.add_module('cls_tower', nn.Sequential(*cls_tower))
        self.add_module('bbox_tower', nn.Sequential(*bbox_tower))
        self.cls_logits = nn.Conv2d(in_channels, 2, kernel_size=3, stride=1, padding=1)
        self.bbox_pred = nn.Conv2d(in_channels, 4, kernel_size=3, stride=1, padding=1)
        self.trans_head = Transformer(cfg)
        for modules in [self.cls_tower, self.bbox_tower, self.trans_head, self.cls_logits, self.bbox_pred]:
            for m in modules.modules():
                if isinstance(m, nn.Conv2d):
                    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                    m.weight.data.normal_(0, math.sqrt(2.0 / n))
                    if m.bias is not None:
                        m.bias.data.zero_()
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1)
                    m.bias.data.zero_()
                elif isinstance(m, nn.Linear):
                    m.weight.data.normal_(0, 0.01)
                    m.bias.data.zero_()

    def forward(self, z_f, x_f):
        zxf_cls = self.trans_head(z_f, x_f)
        zxf_cls = zxf_cls.transpose(2, 3).squeeze(0)
        B, C, WH = zxf_cls.shape
        W, H = int(math.sqrt(WH)), int(math.sqrt(WH))
        zxf_cls = zxf_cls.view(B, C, H, W)
        cls_tower = self.cls_tower(zxf_cls)
        logits = self.cls_logits(cls_tower)
        bbox_tower = self.bbox_tower(zxf_cls)
        bbox_reg = self.bbox_pred(bbox_tower)
        bbox_reg = torch.exp(bbox_reg)
        return logits, bbox_reg


class IOULoss(nn.Module):

    def __init__(self, loc_loss_type):
        super(IOULoss, self).__init__()
        self.loc_loss_type = loc_loss_type

    def forward(self, pred, target, weight=None):
        pred_left = pred[:, 0]
        pred_top = pred[:, 1]
        pred_right = pred[:, 2]
        pred_bottom = pred[:, 3]
        target_left = target[:, 0]
        target_top = target[:, 1]
        target_right = target[:, 2]
        target_bottom = target[:, 3]
        pred_area = (pred_left + pred_right) * (pred_top + pred_bottom)
        target_area = (target_left + target_right) * (target_top + target_bottom)
        w_intersect = torch.min(pred_left, target_left) + torch.min(pred_right, target_right)
        g_w_intersect = torch.max(pred_left, target_left) + torch.max(pred_right, target_right)
        h_intersect = torch.min(pred_bottom, target_bottom) + torch.min(pred_top, target_top)
        g_h_intersect = torch.max(pred_bottom, target_bottom) + torch.max(pred_top, target_top)
        ac_uion = g_w_intersect * g_h_intersect + 1e-07
        area_intersect = w_intersect * h_intersect
        area_union = target_area + pred_area - area_intersect
        ious = (area_intersect + 1.0) / (area_union + 1.0)
        gious = ious - (ac_uion - area_union) / ac_uion
        if self.loc_loss_type == 'iou':
            losses = -torch.log(ious)
        elif self.loc_loss_type == 'linear_iou':
            losses = 1 - ious
        elif self.loc_loss_type == 'giou':
            losses = 1 - gious
        else:
            raise NotImplementedError
        if weight is not None and weight.sum() > 0:
            return (losses * weight).sum() / weight.sum()
        else:
            assert losses.numel() != 0
            return losses.mean()


def mobilenetv3_small(**kwargs):
    """
    Constructs a MobileNetV3-Small model 
    """
    cfgs = [[3, 1, 16, 1, 0, 2], [3, 4.5, 24, 0, 0, 2], [3, 3.67, 24, 0, 0, 1], [5, 4, 40, 1, 1, 2], [5, 6, 40, 1, 1, 1], [5, 6, 40, 1, 1, 1], [5, 3, 48, 1, 1, 1], [5, 3, 48, 1, 1, 1]]
    return MobileNetV3(cfgs, mode='small', **kwargs)


BACKBONES = {'mobilenetv3_small': mobilenetv3_small}


def get_backbone(name, **kwargs):
    return BACKBONES[name](**kwargs)


BANS = {'UPChannelBAN': UPChannelBAN, 'DepthwiseBAN': DepthwiseBAN}


def get_ban_head(name, **kwargs):
    return BANS[name](**kwargs)


class AdjustLayer(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(AdjustLayer, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.downsample = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), nn.BatchNorm2d(out_channels))

    def forward(self, x):
        if self.in_channels != self.out_channels:
            x = self.downsample(x)
        return x


class AdjustAllLayer(nn.Module):

    def __init__(self, in_channels, out_channels):
        super(AdjustAllLayer, self).__init__()
        self.num = len(out_channels)
        if self.num == 1:
            self.downsample = AdjustLayer(in_channels[0], out_channels[0])
        else:
            for i in range(self.num):
                self.add_module('downsample' + str(i + 2), AdjustLayer(in_channels[i], out_channels[i]))

    def forward(self, features):
        if self.num == 1:
            return self.downsample(features)
        else:
            out = []
            for i in range(self.num):
                adj_layer = getattr(self, 'downsample' + str(i + 2))
                out.append(adj_layer(features[i]))
            return out


NECKS = {'AdjustLayer': AdjustLayer, 'AdjustAllLayer': AdjustAllLayer}


def get_neck(name, **kwargs):
    return NECKS[name](**kwargs)


def get_cls_loss(pred, label, select):
    if len(select.size()) == 0 or select.size() == torch.Size([0]):
        return 0
    pred = torch.index_select(pred, 0, select)
    label = torch.index_select(label, 0, select)
    return F.nll_loss(pred, label)


def select_cross_entropy_loss(pred, label):
    pred = pred.view(-1, 2)
    label = label.view(-1)
    pos = label.data.eq(1).nonzero().squeeze()
    neg = label.data.eq(0).nonzero().squeeze()
    loss_pos = get_cls_loss(pred, label, pos)
    loss_neg = get_cls_loss(pred, label, neg)
    return loss_pos * 0.5 + loss_neg * 0.5


linear_iou = IOULoss(loc_loss_type='linear_iou')


def select_iou_loss(pred_loc, label_loc, label_cls):
    label_cls = label_cls.reshape(-1)
    pos = label_cls.data.eq(1).nonzero().squeeze()
    pred_loc = pred_loc.permute(0, 2, 3, 1).reshape(-1, 4)
    pred_loc = torch.index_select(pred_loc, 0, pos)
    label_loc = label_loc.permute(0, 2, 3, 1).reshape(-1, 4)
    label_loc = torch.index_select(label_loc, 0, pos)
    return linear_iou(pred_loc, label_loc)


class ModelBuilder(nn.Module):

    def __init__(self):
        super(ModelBuilder, self).__init__()
        self.backbone = get_backbone(cfg.BACKBONE.TYPE, **cfg.BACKBONE.KWARGS)
        if cfg.ADJUST.ADJUST:
            self.neck = get_neck(cfg.ADJUST.TYPE, **cfg.ADJUST.KWARGS)
        if cfg.BAN.BAN:
            self.ban_head = get_ban_head(cfg.BAN.TYPE, **cfg.BAN.KWARGS)
        self.centerness_loss_func = nn.BCEWithLogitsLoss()

    def compute_centerness_targets(self, reg_targets):
        left_right = reg_targets[:, [0, 2]]
        top_bottom = reg_targets[:, [1, 3]]
        centerness = left_right.min(dim=-1)[0] / left_right.max(dim=-1)[0] * (top_bottom.min(dim=-1)[0] / top_bottom.max(dim=-1)[0])
        return torch.sqrt(centerness)

    def select_cen_loss(self, cen, label_loc, label_cls):
        label_cls = label_cls.reshape(-1)
        pos = label_cls.data.eq(1).nonzero().squeeze()
        cen_preds = cen.view(-1)[pos]
        label_loc = label_loc.permute(0, 2, 3, 1).reshape(-1, 4)
        label_loc = torch.index_select(label_loc, 0, pos)
        cen_targets = self.compute_centerness_targets(label_loc)
        return self.centerness_loss_func(cen_preds, cen_targets)

    def template(self, z):
        zf = self.backbone(z)
        if cfg.ADJUST.ADJUST:
            zf = self.neck(zf)
        self.zf = zf

    def track(self, x):
        xf = self.backbone(x)
        if cfg.ADJUST.ADJUST:
            xf = self.neck(xf)
        cls, loc = self.ban_head(self.zf, xf)
        return {'cls': cls, 'loc': loc}
        return {'cls': cls, 'loc': loc}

    def log_softmax(self, cls):
        if cfg.BAN.BAN:
            cls = cls.permute(0, 2, 3, 1).contiguous()
            cls = F.log_softmax(cls, dim=3)
        return cls

    def forward(self, data):
        """ only used in training
        """
        if len(data) >= 4:
            template = data['template']
            search = data['search']
            label_cls = data['label_cls']
            label_loc = data['label_loc']
            zf = self.backbone(template)
            xf = self.backbone(search)
            if cfg.ADJUST.ADJUST:
                zf = self.neck(zf)
                xf = self.neck(xf)
            cls, loc = self.ban_head(zf, xf)
            cls = self.log_softmax(cls)
            cls_loss = select_cross_entropy_loss(cls, label_cls)
            loc_loss = select_iou_loss(loc, label_loc, label_cls)
            outputs = {}
            outputs['total_loss'] = cfg.TRAIN.CLS_WEIGHT * cls_loss + cfg.TRAIN.LOC_WEIGHT * loc_loss
            outputs['cls_loss'] = cls_loss
            outputs['loc_loss'] = loc_loss
            return outputs
        else:
            xf = self.backbone(data)
            if cfg.ADJUST.ADJUST:
                xf = self.neck(xf)
            cls, loc = self.ban_head(self.zf, xf)
            return {'cls': cls, 'loc': loc}


inited = False


def get_world_size():
    if not inited:
        raise Exception('dist not inited')
    return world_size


def broadcast_buffers(model, method=0):
    """ broadcast model buffers """
    if method == 0:
        return
    world_size = get_world_size()
    for b in model._all_buffers():
        if method == 1:
            dist.broadcast(b, 0)
        elif method == 2:
            dist.all_reduce(b)
            b /= world_size
        else:
            raise Exception('Invalid buffer broadcast code {}'.format(method))


def broadcast_params(model):
    """ broadcast model parameters """
    for p in model.state_dict().values():
        dist.broadcast(p, 0)


class DistModule(nn.Module):

    def __init__(self, module, bn_method=0):
        super(DistModule, self).__init__()
        self.module = module
        self.bn_method = bn_method
        if get_world_size() > 1:
            broadcast_params(self.module)
        else:
            self.bn_method = 0

    def forward(self, *args, **kwargs):
        broadcast_buffers(self.module, self.bn_method)
        return self.module(*args, **kwargs)

    def train(self, mode=True):
        super(DistModule, self).train(mode)
        self.module.train(mode)
        return self


def conv_bn(inp, oup, stride, padding=1):
    return nn.Sequential(nn.Conv2d(inp, oup, 3, stride, padding, bias=False), nn.BatchNorm2d(oup), nn.ReLU6(inplace=True))


class MobileNetV2(nn.Sequential):

    def __init__(self, width_mult=1.0, used_layers=[3, 5, 7]):
        super(MobileNetV2, self).__init__()
        self.interverted_residual_setting = [[1, 16, 1, 1, 1], [6, 24, 2, 2, 1], [6, 32, 3, 2, 1], [6, 64, 4, 1, 2], [6, 96, 3, 1, 2], [6, 160, 3, 1, 4], [6, 320, 1, 1, 4]]
        self.channels = [24, 32, 96, 320]
        self.channels = [int(c * width_mult) for c in self.channels]
        input_channel = int(32 * width_mult)
        self.last_channel = int(1280 * width_mult) if width_mult > 1.0 else 1280
        self.add_module('layer0', conv_bn(3, input_channel, 2, 0))
        last_dilation = 1
        self.used_layers = used_layers
        for idx, (t, c, n, s, d) in enumerate(self.interverted_residual_setting, start=1):
            output_channel = int(c * width_mult)
            layers = []
            for i in range(n):
                if i == 0:
                    if d == last_dilation:
                        dd = d
                    else:
                        dd = max(d // 2, 1)
                    layers.append(InvertedResidual(input_channel, output_channel, s, t, dd))
                else:
                    layers.append(InvertedResidual(input_channel, output_channel, 1, t, d))
                input_channel = output_channel
            last_dilation = d
            self.add_module('layer%d' % idx, nn.Sequential(*layers))

    def forward(self, x):
        outputs = []
        for idx in range(8):
            name = 'layer%d' % idx
            x = getattr(self, name)(x)
            outputs.append(x)
        p0, p1, p2, p3, p4 = [outputs[i] for i in [1, 2, 3, 5, 7]]
        out = [outputs[i] for i in self.used_layers]
        return out


def conv3x3(in_planes, out_planes, stride=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, bias=False, dilation=dilation)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
        super(BasicBlock, self).__init__()
        padding = 2 - stride
        if dilation > 1:
            padding = dilation
        dd = dilation
        pad = padding
        if downsample is not None and dilation > 1:
            dd = dilation // 2
            pad = dd
        self.conv1 = nn.Conv2d(inplanes, planes, stride=stride, dilation=dd, bias=False, kernel_size=3, padding=pad)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes, dilation=dilation)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out


class Bottleneck(nn.Module):

    def __init__(self, in_channels: int=96, mid_channels: int=96, dilation: int=1):
        super(Bottleneck, self).__init__()
        self.in_channels = in_channels
        self.mid_channels = mid_channels
        self.pw1 = nn.Sequential(nn.Conv2d(in_channels, mid_channels, kernel_size=1, padding=0, bias=False), nn.BatchNorm2d(mid_channels), nn.ReLU6(inplace=True))
        self.dw = nn.Sequential(nn.Conv2d(mid_channels, mid_channels, kernel_size=3, padding=1, groups=mid_channels, bias=False), nn.BatchNorm2d(mid_channels), nn.ReLU6(inplace=True))
        self.pw2 = nn.Sequential(nn.Conv2d(mid_channels, in_channels, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(in_channels))

    def forward(self, x):
        identity = x
        if self.mid_channels != self.in_channels:
            out = self.pw1(x)
            out = self.dw(out)
            out = self.pw2(out)
        else:
            out = self.dw(x)
            out = self.pw2(out)
        out = out + identity
        return out


class ResNet(nn.Module):

    def __init__(self, block, layers, used_layers):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.feature_size = 128 * block.expansion
        self.used_layers = used_layers
        layer3 = True if 3 in used_layers or 4 in used_layers else False
        layer4 = True if 4 in used_layers else False
        if layer3:
            self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)
            self.feature_size = (256 + 128) * block.expansion
        else:
            self.layer3 = lambda x: x
        if layer4:
            self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4)
            self.feature_size = 512 * block.expansion
        else:
            self.layer4 = lambda x: x
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2.0 / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
        downsample = None
        dd = dilation
        if stride != 1 or self.inplanes != planes * block.expansion:
            if stride == 1 and dilation == 1:
                downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion))
            else:
                if dilation > 1:
                    dd = dilation // 2
                    padding = dd
                else:
                    dd = 1
                    padding = 0
                downsample = nn.Sequential(nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=3, stride=stride, bias=False, padding=padding, dilation=dd), nn.BatchNorm2d(planes * block.expansion))
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, dilation=dilation))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, dilation=dilation))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x_ = self.relu(x)
        x = self.maxpool(x_)
        p1 = self.layer1(x)
        p2 = self.layer2(p1)
        p3 = self.layer3(p2)
        p4 = self.layer4(p3)
        out = [x_, p1, p2, p3, p4]
        out = [out[i] for i in self.used_layers]
        if len(out) == 1:
            return out[0]
        else:
            return out


class MultiBAN(BAN):

    def __init__(self, in_channels, cls_out_channels, weighted=False):
        super(MultiBAN, self).__init__()
        self.weighted = weighted
        for i in range(len(in_channels)):
            self.add_module('box' + str(i + 2), DepthwiseBAN(in_channels[i], in_channels[i], cls_out_channels))
        if self.weighted:
            self.cls_weight = nn.Parameter(torch.ones(len(in_channels)))
            self.loc_weight = nn.Parameter(torch.ones(len(in_channels)))
        self.loc_scale = nn.Parameter(torch.ones(len(in_channels)))

    def forward(self, z_fs, x_fs):
        cls = []
        loc = []
        for idx, (z_f, x_f) in enumerate(zip(z_fs, x_fs), start=2):
            box = getattr(self, 'box' + str(idx))
            c, l = box(z_f, x_f)
            cls.append(c)
            loc.append(torch.exp(l * self.loc_scale[idx - 2]))
        if self.weighted:
            cls_weight = F.softmax(self.cls_weight, 0)
            loc_weight = F.softmax(self.loc_weight, 0)

        def avg(lst):
            return sum(lst) / len(lst)

        def weighted_avg(lst, weight):
            s = 0
            for i in range(len(weight)):
                s += lst[i] * weight[i]
            return s
        if self.weighted:
            return weighted_avg(cls, cls_weight), weighted_avg(loc, loc_weight)
        else:
            return avg(cls), avg(loc)


class CARHead(torch.nn.Module):

    def __init__(self, cfg, in_channels):
        """
        Arguments:
            in_channels (int): number of channels of the input feature
        """
        super(CARHead, self).__init__()
        num_classes = cfg.TRAIN.NUM_CLASSES
        cls_tower = []
        bbox_tower = []
        for i in range(cfg.TRAIN.NUM_CONVS):
            cls_tower.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1))
            cls_tower.append(nn.GroupNorm(32, in_channels))
            cls_tower.append(nn.ReLU())
            bbox_tower.append(nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1))
            bbox_tower.append(nn.GroupNorm(32, in_channels))
            bbox_tower.append(nn.ReLU())
        self.add_module('cls_tower', nn.Sequential(*cls_tower))
        self.add_module('bbox_tower', nn.Sequential(*bbox_tower))
        self.cls_logits = nn.Conv2d(in_channels, num_classes, kernel_size=3, stride=1, padding=1)
        self.bbox_pred = nn.Conv2d(in_channels, 4, kernel_size=3, stride=1, padding=1)
        self.centerness = nn.Conv2d(in_channels, 1, kernel_size=3, stride=1, padding=1)
        out_channels = in_channels
        self.cls_dw = DepthwiseXCorr(in_channels, out_channels)
        self.reg_dw = DepthwiseXCorr(in_channels, out_channels)
        for modules in [self.cls_tower, self.bbox_tower, self.cls_logits, self.bbox_pred, self.centerness]:
            for l in modules.modules():
                if isinstance(l, nn.Conv2d):
                    torch.nn.init.normal_(l.weight, std=0.01)
                    torch.nn.init.constant_(l.bias, 0)
        prior_prob = cfg.TRAIN.PRIOR_PROB
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        torch.nn.init.constant_(self.cls_logits.bias, bias_value)

    def forward(self, z_f, x_f):
        x_cls = self.cls_dw(z_f, x_f)
        x_reg = self.reg_dw(z_f, x_f)
        cls_tower = self.cls_tower(x_cls)
        logits = self.cls_logits(cls_tower)
        centerness = self.centerness(cls_tower)
        bbox_tower = self.bbox_tower(x_reg)
        bbox_reg = self.bbox_pred(bbox_tower)
        bbox_reg = torch.exp(bbox_reg)
        return logits, bbox_reg, centerness


def bbox_area(bbox):
    return (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1])


def clip_box(bbox, clip_box, alpha):
    """Clip the bounding boxes to the borders of an image
    
    Parameters
    ----------
    
    bbox: numpy.ndarray
        Numpy array containing bounding boxes of shape `N X 4` where N is the 
        number of bounding boxes and the bounding boxes are represented in the
        format `x1 y1 x2 y2`
    
    clip_box: numpy.ndarray
        An array of shape (4,) specifying the diagonal co-ordinates of the image
        The coordinates are represented in the format `x1 y1 x2 y2`
        
    alpha: float
        If the fraction of a bounding box left in the image after being clipped is 
        less than `alpha` the bounding box is dropped. 
    
    Returns
    -------
    
    numpy.ndarray
        Numpy array containing **clipped** bounding boxes of shape `N X 4` where N is the 
        number of bounding boxes left are being clipped and the bounding boxes are represented in the
        format `x1 y1 x2 y2` 
    
    """
    ar_ = bbox_area(bbox)
    x_min = np.maximum(bbox[:, 0], clip_box[0]).reshape(-1, 1)
    y_min = np.maximum(bbox[:, 1], clip_box[1]).reshape(-1, 1)
    x_max = np.minimum(bbox[:, 2], clip_box[2]).reshape(-1, 1)
    y_max = np.minimum(bbox[:, 3], clip_box[3]).reshape(-1, 1)
    bbox = np.hstack((x_min, y_min, x_max, y_max, bbox[:, 4:]))
    delta_area = (ar_ - bbox_area(bbox)) / ar_
    mask = (delta_area < 1 - alpha).astype(int)
    bbox = bbox[mask == 1, :]
    return bbox


class Scale(object):
    """Scales the image    
    Bounding boxes which have an area of less than 25% in the remaining in the 
    transformed image is dropped. The resolution is maintained, and the remaining
    area if any is filled by black color
    Parameters
    ----------
    scale_x: float
        The factor by which the image is scaled horizontally
        
    scale_y: float
        The factor by which the image is scaled vertically
    Returns
    -------
    numpy.ndaaray
        Scaled image in the numpy format of shape `HxWxC`
    
    numpy.ndarray
        Tranformed bounding box co-ordinates of the format `n x 4` where n is 
        number of bounding boxes and 4 represents `x1,y1,x2,y2` of the box  
    """

    def __init__(self, scale_x=0.2, scale_y=0.2):
        self.scale_x = scale_x
        self.scale_y = scale_y

    def __call__(self, img, bboxes):
        img_shape = img.shape
        resize_scale_x = 1 + self.scale_x
        resize_scale_y = 1 + self.scale_y
        img = cv2.resize(img, None, fx=resize_scale_x, fy=resize_scale_y)
        bboxes[:, :4] *= [resize_scale_x, resize_scale_y, resize_scale_x, resize_scale_y]
        canvas = np.zeros(img_shape, dtype=np.uint8)
        y_lim = int(min(resize_scale_y, 1) * img_shape[0])
        x_lim = int(min(resize_scale_x, 1) * img_shape[1])
        canvas[:y_lim, :x_lim, :] = img[:y_lim, :x_lim, :]
        img = canvas
        bboxes = clip_box(bboxes, [0, 0, 1 + img_shape[1], img_shape[0]], 0.25)
        return img, bboxes


class _BatchNorm2d(nn.BatchNorm2d):

    def __init__(self, num_features, *args, **kwargs):
        super(_BatchNorm2d, self).__init__(num_features, *args, eps=1e-06, momentum=0.05, **kwargs)


class _AlexNet(nn.Module):

    def forward(self, x):
        x = self.features(x)
        return x


class AlexNetV0(_AlexNet):
    output_stride = 8

    def __init__(self):
        super(AlexNetV0, self).__init__()
        self.features = nn.Sequential(nn.Conv2d(3, 96, 11, 2), _BatchNorm2d(96), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2), nn.Conv2d(96, 256, 5, 1), _BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2), nn.Conv2d(256, 384, 3, 1), _BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 384, 3, 1), _BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 256, 3, 1))


class AlexNetV1(_AlexNet):
    output_stride = 8

    def __init__(self):
        super(AlexNetV1, self).__init__()
        self.features = nn.Sequential(nn.Conv2d(3, 96, 11, 2), _BatchNorm2d(96), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2), nn.Conv2d(96, 256, 5, 1, groups=2), _BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2), nn.Conv2d(256, 384, 3, 1), _BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 384, 3, 1, groups=2), _BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 256, 3, 1, groups=2))


class AlexNetV2(_AlexNet):
    output_stride = 4

    def __init__(self):
        super(AlexNetV2, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(3, 96, 11, 2), _BatchNorm2d(96), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2))
        self.conv2 = nn.Sequential(nn.Conv2d(96, 256, 5, 1, groups=2), _BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(3, 1))
        self.conv3 = nn.Sequential(nn.Conv2d(256, 384, 3, 1), _BatchNorm2d(384), nn.ReLU(inplace=True))
        self.conv4 = nn.Sequential(nn.Conv2d(384, 384, 3, 1, groups=2), _BatchNorm2d(384), nn.ReLU(inplace=True))
        self.conv5 = nn.Sequential(nn.Conv2d(384, 256, 3, 1, groups=2))


class AlexNetV3(_AlexNet):
    output_stride = 8

    def __init__(self):
        super(AlexNetV3, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(3, 192, 11, 2), _BatchNorm2d(192), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2))
        self.conv2 = nn.Sequential(nn.Conv2d(192, 512, 5, 1), _BatchNorm2d(512), nn.ReLU(inplace=True), nn.MaxPool2d(3, 2))
        self.conv3 = nn.Sequential(nn.Conv2d(512, 768, 3, 1), _BatchNorm2d(768), nn.ReLU(inplace=True))
        self.conv4 = nn.Sequential(nn.Conv2d(768, 768, 3, 1), _BatchNorm2d(768), nn.ReLU(inplace=True))
        self.conv5 = nn.Sequential(nn.Conv2d(768, 512, 3, 1), _BatchNorm2d(512))


def center_crop(x):
    """
    center crop layer. crop [1:-2] to eliminate padding influence.
    Crop 1 element around the tensor
    input x can be a Variable or Tensor
    """
    return x[:, :, 1:-1, 1:-1].contiguous()


eps = 1e-07


class Bottleneck_CI(nn.Module):
    """
    Bottleneck with center crop layer, utilized in CVPR2019 model
    """
    expansion = 4

    def __init__(self, inplanes, planes, last_relu, stride=1, downsample=None, dilation=1):
        super(Bottleneck_CI, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        padding = 1
        if abs(dilation - 2) < eps:
            padding = 2
        if abs(dilation - 3) < eps:
            padding = 3
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=padding, bias=False, dilation=dilation)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
        self.last_relu = last_relu

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        if self.last_relu:
            out = self.relu(out)
        out = center_crop(out)
        return out


class ResNet22(nn.Module):
    """
    default: unfix gradually (lr: 1r-2 ~ 1e-5)
    optional: unfix all at first with small lr (lr: 1e-7 ~ 1e-3)
    """

    def __init__(self):
        super(ResNet22, self).__init__()
        self.features = ResNet(Bottleneck_CI, [3, 4], [True, False], [False, True])
        self.feature_size = 512
        self.train_num = 0
        self.unfix(0.0)

    def forward(self, x):
        x = self.features(x)
        return x

    def unfix(self, ratio):
        """
        unfix gradually as paper said
        """
        if abs(ratio - 0.0) < eps:
            self.train_num = 2
            self.unlock()
            return True
        elif abs(ratio - 0.1) < eps:
            self.train_num = 3
            self.unlock()
            return True
        elif abs(ratio - 0.2) < eps:
            self.train_num = 4
            self.unlock()
            return True
        elif abs(ratio - 0.3) < eps:
            self.train_num = 6
            self.unlock()
            return True
        elif abs(ratio - 0.5) < eps:
            self.train_num = 7
            self.unlock()
            return True
        elif abs(ratio - 0.6) < eps:
            self.train_num = 8
            self.unlock()
            return True
        elif abs(ratio - 0.7) < eps:
            self.train_num = 9
            self.unlock()
            return True
        return False

    def unlock(self):
        for p in self.parameters():
            p.requires_grad = False
        for i in range(1, self.train_num):
            if i <= 5:
                m = self.features.layer2[-i]
            elif i <= 8:
                m = self.features.layer1[-(i - 5)]
            else:
                m = self.features
            for p in m.parameters():
                p.requires_grad = True
        self.eval()
        self.train()

    def train(self, mode=True):
        self.training = mode
        if mode == False:
            super(ResNet22, self).train(False)
        else:
            for i in range(self.train_num):
                if i <= 5:
                    m = self.features.layer2[-i]
                elif i <= 8:
                    m = self.features.layer1[-(i - 5)]
                else:
                    m = self.features
                m.train(mode)
        return self


class Inception(nn.Module):
    """
    Inception with 22 layer utilized in CVPR2019 paper.
    Usage: Inception(InceptionM, [3, 4], [True, False])
    """

    def __init__(self, block, layers):
        self.inplanes = 64
        super(Inception, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.layer1 = self._make_layer(block, 64, 64, layers[0], pool=False)
        self.layer2 = self._make_layer(block, 320, 128, layers[1], pool=True, last_relu=False)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal(m.weight, mode='fan_out')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant(m.weight, 1)
                nn.init.constant(m.bias, 0)

    def _make_layer(self, block, inchannels, planes, blocks, pool=True, last_relu=True):
        layers = []
        for i in range(0, blocks):
            if i == 0:
                self.inchannels = inchannels
            else:
                self.inchannels = planes * 5
            if i == 1 and pool:
                layers.append(self.maxpool)
            if i == blocks - 1 and not last_relu:
                layers.append(block(self.inchannels, planes, last_relu))
            else:
                layers.append(block(self.inchannels, planes))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool1(x)
        x = self.layer1(x)
        x = self.layer2(x)
        return x


class BasicConv2d_1x1(nn.Module):
    """
    1*1 branch of inception
    """

    def __init__(self, in_channels, out_channels, last_relu=True, **kwargs):
        super(BasicConv2d_1x1, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
        self.last_relu = last_relu

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.last_relu:
            return F.relu(x, inplace=True)
        else:
            return x


class BasicConv2d_3x3(nn.Module):
    """
    3*3 branch of inception
    """
    expansion = 4

    def __init__(self, inplanes, planes, last_relu=True):
        super(BasicConv2d_3x3, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.last_relu = last_relu

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.last_relu:
            out = self.relu(out)
        return out


class InceptionM(nn.Module):
    """
    Inception module with 1*1 and 3*3 branch
    """

    def __init__(self, in_channels, planes, last_relu=True):
        super(InceptionM, self).__init__()
        self.branch3x3 = BasicConv2d_3x3(in_channels, planes, last_relu)
        self.branch1x1 = BasicConv2d_1x1(in_channels, planes, last_relu, kernel_size=1)

    def forward(self, x):
        branch3x3 = self.branch3x3(x)
        branch1x1 = self.branch1x1(x)
        outputs = [branch3x3, branch1x1]
        return center_crop(torch.cat(outputs, 1))


class Incep22(nn.Module):

    def __init__(self):
        super(Incep22, self).__init__()
        self.features = Inception(InceptionM, [3, 4])
        self.feature_size = 640

    def forward(self, x):
        x = self.features(x)
        return x


class BasicBlock_C(nn.Module):
    """
    increasing cardinality is a more effective way of
    gaining accuracy than going deeper or wider
    """

    def __init__(self, in_planes, bottleneck_width=4, cardinality=32, expansion=2, last_relu=True):
        super(BasicBlock_C, self).__init__()
        inner_width = cardinality * bottleneck_width
        self.expansion = expansion
        self.basic = nn.Sequential(OrderedDict([('conv1_0', nn.Conv2d(in_planes, inner_width, 1, stride=1, bias=False)), ('bn1', nn.BatchNorm2d(inner_width)), ('act0', nn.ReLU()), ('conv3_0', nn.Conv2d(inner_width, inner_width, 3, stride=1, padding=1, groups=cardinality, bias=False)), ('bn2', nn.BatchNorm2d(inner_width)), ('act1', nn.ReLU()), ('conv1_1', nn.Conv2d(inner_width, inner_width * self.expansion, 1, stride=1, bias=False)), ('bn3', nn.BatchNorm2d(inner_width * self.expansion))]))
        self.shortcut = nn.Sequential()
        if in_planes != inner_width * self.expansion:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, inner_width * self.expansion, 1, stride=1, bias=False))
        self.bn0 = nn.BatchNorm2d(self.expansion * inner_width)
        self.last_relu = last_relu

    def forward(self, x):
        out = self.basic(x)
        out += self.shortcut(x)
        if self.last_relu:
            return center_crop(F.relu(self.bn0(out)))
        else:
            return center_crop(self.bn0(out))


class ResNeXt(nn.Module):
    """
    ResNeXt with 22 layer utilized in CVPR2019 paper.
    Usage: ResNeXt([3, 4], 32, 4)
    """

    def __init__(self, num_blocks, cardinality, bottleneck_width, expansion=2):
        super(ResNeXt, self).__init__()
        self.cardinality = cardinality
        self.bottleneck_width = bottleneck_width
        self.in_planes = 64
        self.expansion = expansion
        self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv0 = nn.Conv2d(3, self.in_planes, kernel_size=7, stride=2, padding=0)
        self.bn0 = nn.BatchNorm2d(self.in_planes)
        self.layer1 = self._make_layer(num_blocks[0], last_relu=True)
        self.layer2 = self._make_layer(num_blocks[1], last_relu=False, stride2pool=True)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal(m.weight, mode='fan_out')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant(m.weight, 1)
                nn.init.constant(m.bias, 0)

    def _make_layer(self, num_blocks, last_relu=True, stride2pool=False):
        layers = []
        for i in range(0, num_blocks):
            if i == num_blocks - 1:
                layers.append(BasicBlock_C(self.in_planes, self.bottleneck_width, self.cardinality, self.expansion, last_relu=last_relu))
            else:
                layers.append(BasicBlock_C(self.in_planes, self.bottleneck_width, self.cardinality, self.expansion))
            self.in_planes = self.expansion * self.bottleneck_width * self.cardinality
            if i == 0 and stride2pool:
                layers.append(self.maxpool)
        self.bottleneck_width *= 2
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn0(self.conv0(x)))
        out = self.maxpool1(out)
        out = self.layer1(out)
        out = self.layer2(out)
        return out


class ResNeXt22(nn.Module):

    def __init__(self):
        super(ResNeXt22, self).__init__()
        self.features = ResNeXt(num_blocks=[3, 4], cardinality=32, bottleneck_width=4)
        self.feature_size = 512

    def forward(self, x):
        x = self.features(x)
        return x


class Bottleneck_BIG_CI(nn.Module):
    """
    Bottleneck with center crop layer, double channels in 3*3 conv layer in shortcut branch
    """
    expansion = 4

    def __init__(self, inplanes, planes, last_relu, stride=1, downsample=None, dilation=1):
        super(Bottleneck_BIG_CI, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        padding = 1
        if abs(dilation - 2) < eps:
            padding = 2
        if abs(dilation - 3) < eps:
            padding = 3
        self.conv2 = nn.Conv2d(planes, planes * 2, kernel_size=3, stride=stride, padding=padding, bias=False, dilation=dilation)
        self.bn2 = nn.BatchNorm2d(planes * 2)
        self.conv3 = nn.Conv2d(planes * 2, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride
        self.last_relu = last_relu

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            residual = self.downsample(x)
        out += residual
        if self.last_relu:
            out = self.relu(out)
        out = center_crop(out)
        return out


class ResNet22W(nn.Module):
    """
    ResNet22W: double 3*3 layer (only) channels in residual blob
    """

    def __init__(self):
        super(ResNet22W, self).__init__()
        self.features = ResNet(Bottleneck_BIG_CI, [3, 4], [True, False], [False, True], firstchannels=64, channels=[64, 128])
        self.feature_size = 512

    def forward(self, x):
        x = self.features(x)
        return x


class SiamFC(nn.Module):

    def __init__(self, out_scale=0.001):
        super(SiamFC, self).__init__()
        self.out_scale = out_scale

    def forward(self, z, x):
        return self._fast_xcorr(z, x) * self.out_scale

    def _fast_xcorr(self, z, x):
        nz = z.size(0)
        nx, c, h, w = x.size()
        x = x.view(-1, nz * c, h, w)
        out = F.conv2d(x, z, groups=nz)
        out = out.view(nx, -1, out.size(-2), out.size(-1))
        return out


class BalancedLoss(nn.Module):

    def __init__(self, neg_weight=1.0):
        super(BalancedLoss, self).__init__()
        self.neg_weight = neg_weight

    def forward(self, input, target):
        pos_mask = target == 1
        neg_mask = target == 0
        pos_num = pos_mask.sum().float()
        neg_num = neg_mask.sum().float()
        weight = target.new_zeros(target.size())
        weight[pos_mask] = 1 / pos_num
        weight[neg_mask] = 1 / neg_num * self.neg_weight
        weight /= weight.sum()
        return F.binary_cross_entropy_with_logits(input, target, weight, reduction='sum')


def log_minus_sigmoid(x):
    return torch.clamp(-x, max=0) - torch.log(1 + torch.exp(-torch.abs(x))) + 0.5 * torch.clamp(x, min=0, max=0)


def log_sigmoid(x):
    return torch.clamp(x, max=0) - torch.log(1 + torch.exp(-torch.abs(x))) + 0.5 * torch.clamp(x, min=0, max=0)


class FocalLoss(nn.Module):

    def __init__(self, gamma=2):
        super(FocalLoss, self).__init__()
        self.gamma = gamma

    def forward(self, input, target):
        pos_log_sig = log_sigmoid(input)
        neg_log_sig = log_minus_sigmoid(input)
        prob = torch.sigmoid(input)
        pos_weight = torch.pow(1 - prob, self.gamma)
        neg_weight = torch.pow(prob, self.gamma)
        loss = -(target * pos_weight * pos_log_sig + (1 - target) * neg_weight * neg_log_sig)
        avg_weight = target * pos_weight + (1 - target) * neg_weight
        loss /= avg_weight.mean()
        return loss.mean()


class GHMCLoss(nn.Module):

    def __init__(self, bins=30, momentum=0.5):
        super(GHMCLoss, self).__init__()
        self.bins = bins
        self.momentum = momentum
        self.edges = [(t / bins) for t in range(bins + 1)]
        self.edges[-1] += 1e-06
        if momentum > 0:
            self.acc_sum = [(0.0) for _ in range(bins)]

    def forward(self, input, target):
        edges = self.edges
        mmt = self.momentum
        weights = torch.zeros_like(input)
        g = torch.abs(input.sigmoid().detach() - target)
        tot = input.numel()
        n = 0
        for i in range(self.bins):
            inds = (g >= edges[i]) & (g < edges[i + 1])
            num_in_bin = inds.sum().item()
            if num_in_bin > 0:
                if mmt > 0:
                    self.acc_sum[i] = mmt * self.acc_sum[i] + (1 - mmt) * num_in_bin
                    weights[inds] = tot / self.acc_sum[i]
                else:
                    weights[inds] = tot / num_in_bin
                n += 1
        if n > 0:
            weights /= weights.mean()
        loss = F.binary_cross_entropy_with_logits(input, target, weights, reduction='sum') / tot
        return loss


class OHNMLoss(nn.Module):

    def __init__(self, neg_ratio=3.0):
        super(OHNMLoss, self).__init__()
        self.neg_ratio = neg_ratio

    def forward(self, input, target):
        pos_logits = input[target > 0]
        pos_labels = target[target > 0]
        neg_logits = input[target == 0]
        neg_labels = target[target == 0]
        pos_num = pos_logits.numel()
        neg_num = int(pos_num * self.neg_ratio)
        neg_logits, neg_indices = neg_logits.topk(neg_num)
        neg_labels = neg_labels[neg_indices]
        loss = F.binary_cross_entropy_with_logits(torch.cat([pos_logits, neg_logits]), torch.cat([pos_labels, neg_labels]), reduction='mean')
        return loss


class SiamFCNet(nn.Module):

    def __init__(self, backbone, head):
        super(SiamFCNet, self).__init__()
        self.features = backbone
        self.head = head

    def forward(self, z, x):
        z = self.features(z)
        x = self.features(x)
        return self.head(z, x)


class SiamRPNNet(nn.Module):

    def __init__(self):
        super(SiamRPNNet, self).__init__()
        self.featureExtract = nn.Sequential(nn.Conv2d(3, 96, 11, stride=2), nn.BatchNorm2d(96), nn.MaxPool2d(3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(96, 256, 5), nn.BatchNorm2d(256), nn.MaxPool2d(3, stride=2), nn.ReLU(inplace=True), nn.Conv2d(256, 384, 3), nn.BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 384, 3), nn.BatchNorm2d(384), nn.ReLU(inplace=True), nn.Conv2d(384, 256, 3), nn.BatchNorm2d(256))
        self.anchor_num = config.anchor_num
        self.input_size = config.instance_size
        self.score_displacement = int((self.input_size - config.exemplar_size) / config.total_stride)
        self.conv_cls1 = nn.Conv2d(256, 256 * 2 * self.anchor_num, kernel_size=3, stride=1, padding=0)
        self.conv_r1 = nn.Conv2d(256, 256 * 4 * self.anchor_num, kernel_size=3, stride=1, padding=0)
        self.conv_cls2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=0)
        self.conv_r2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=0)
        self.regress_adjust = nn.Conv2d(4 * self.anchor_num, 4 * self.anchor_num, 1)

    def forward(self, template, detection):
        N = template.size(0)
        template_feature = self.featureExtract(template)
        detection_feature = self.featureExtract(detection)
        kernel_score = self.conv_cls1(template_feature).view(N, 2 * self.anchor_num, 256, 4, 4)
        kernel_regression = self.conv_r1(template_feature).view(N, 4 * self.anchor_num, 256, 4, 4)
        conv_score = self.conv_cls2(detection_feature)
        conv_regression = self.conv_r2(detection_feature)
        conv_scores = conv_score.reshape(1, -1, self.score_displacement + 4, self.score_displacement + 4)
        score_filters = kernel_score.reshape(-1, 256, 4, 4)
        pred_score = F.conv2d(conv_scores, score_filters, groups=N).reshape(N, 10, self.score_displacement + 1, self.score_displacement + 1)
        conv_reg = conv_regression.reshape(1, -1, self.score_displacement + 4, self.score_displacement + 4)
        reg_filters = kernel_regression.reshape(-1, 256, 4, 4)
        pred_regression = self.regress_adjust(F.conv2d(conv_reg, reg_filters, groups=N).reshape(N, 20, self.score_displacement + 1, self.score_displacement + 1))
        return pred_score, pred_regression

    def track_init(self, template):
        N = template.size(0)
        template_feature = self.featureExtract(template)
        kernel_score = self.conv_cls1(template_feature).view(N, 2 * self.anchor_num, 256, 4, 4)
        kernel_regression = self.conv_r1(template_feature).view(N, 4 * self.anchor_num, 256, 4, 4)
        self.score_filters = kernel_score.reshape(-1, 256, 4, 4)
        self.reg_filters = kernel_regression.reshape(-1, 256, 4, 4)

    def track(self, detection):
        N = detection.size(0)
        detection_feature = self.featureExtract(detection)
        conv_score = self.conv_cls2(detection_feature)
        conv_regression = self.conv_r2(detection_feature)
        conv_scores = conv_score.reshape(1, -1, self.score_displacement + 4, self.score_displacement + 4)
        pred_score = F.conv2d(conv_scores, self.score_filters, groups=N).reshape(N, 10, self.score_displacement + 1, self.score_displacement + 1)
        conv_reg = conv_regression.reshape(1, -1, self.score_displacement + 4, self.score_displacement + 4)
        pred_regression = self.regress_adjust(F.conv2d(conv_reg, self.reg_filters, groups=N).reshape(N, 20, self.score_displacement + 1, self.score_displacement + 1))
        return pred_score, pred_regression


class BasicConv2d(nn.Module):

    def __init__(self, in_channels, out_channels, **kwargs):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return F.relu(x, inplace=True)


class InceptionA(nn.Module):

    def __init__(self, in_channels, pool_features):
        super(InceptionA, self).__init__()
        self.branch1x1 = BasicConv2d(in_channels, 64, kernel_size=1)
        self.branch5x5_1 = BasicConv2d(in_channels, 48, kernel_size=1)
        self.branch5x5_2 = BasicConv2d(48, 64, kernel_size=5, padding=2)
        self.branch3x3dbl_1 = BasicConv2d(in_channels, 64, kernel_size=1)
        self.branch3x3dbl_2 = BasicConv2d(64, 96, kernel_size=3, padding=1)
        self.branch3x3dbl_3 = BasicConv2d(96, 96, kernel_size=3, padding=1)
        self.branch_pool = BasicConv2d(in_channels, pool_features, kernel_size=1)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)
        branch5x5 = self.branch5x5_1(x)
        branch5x5 = self.branch5x5_2(branch5x5)
        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)
        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


class InceptionB(nn.Module):

    def __init__(self, in_channels):
        super(InceptionB, self).__init__()
        self.branch3x3 = BasicConv2d(in_channels, 384, kernel_size=3, stride=2)
        self.branch3x3dbl_1 = BasicConv2d(in_channels, 64, kernel_size=1)
        self.branch3x3dbl_2 = BasicConv2d(64, 96, kernel_size=3, padding=1)
        self.branch3x3dbl_3 = BasicConv2d(96, 96, kernel_size=3, stride=2)

    def forward(self, x):
        branch3x3 = self.branch3x3(x)
        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
        outputs = [branch3x3, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


class InceptionC(nn.Module):

    def __init__(self, in_channels, channels_7x7):
        super(InceptionC, self).__init__()
        self.branch1x1 = BasicConv2d(in_channels, 192, kernel_size=1)
        c7 = channels_7x7
        self.branch7x7_1 = BasicConv2d(in_channels, c7, kernel_size=1)
        self.branch7x7_2 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3))
        self.branch7x7_3 = BasicConv2d(c7, 192, kernel_size=(7, 1), padding=(3, 0))
        self.branch7x7dbl_1 = BasicConv2d(in_channels, c7, kernel_size=1)
        self.branch7x7dbl_2 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0))
        self.branch7x7dbl_3 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3))
        self.branch7x7dbl_4 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0))
        self.branch7x7dbl_5 = BasicConv2d(c7, 192, kernel_size=(1, 7), padding=(0, 3))
        self.branch_pool = BasicConv2d(in_channels, 192, kernel_size=1)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)
        branch7x7 = self.branch7x7_1(x)
        branch7x7 = self.branch7x7_2(branch7x7)
        branch7x7 = self.branch7x7_3(branch7x7)
        branch7x7dbl = self.branch7x7dbl_1(x)
        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)
        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
        return torch.cat(outputs, 1)


class InceptionD(nn.Module):

    def __init__(self, in_channels):
        super(InceptionD, self).__init__()
        self.branch3x3_1 = BasicConv2d(in_channels, 192, kernel_size=1)
        self.branch3x3_2 = BasicConv2d(192, 320, kernel_size=3, stride=2)
        self.branch7x7x3_1 = BasicConv2d(in_channels, 192, kernel_size=1)
        self.branch7x7x3_2 = BasicConv2d(192, 192, kernel_size=(1, 7), padding=(0, 3))
        self.branch7x7x3_3 = BasicConv2d(192, 192, kernel_size=(7, 1), padding=(3, 0))
        self.branch7x7x3_4 = BasicConv2d(192, 192, kernel_size=3, stride=2)

    def forward(self, x):
        branch3x3 = self.branch3x3_1(x)
        branch3x3 = self.branch3x3_2(branch3x3)
        branch7x7x3 = self.branch7x7x3_1(x)
        branch7x7x3 = self.branch7x7x3_2(branch7x7x3)
        branch7x7x3 = self.branch7x7x3_3(branch7x7x3)
        branch7x7x3 = self.branch7x7x3_4(branch7x7x3)
        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
        outputs = [branch3x3, branch7x7x3, branch_pool]
        return torch.cat(outputs, 1)


class InceptionE(nn.Module):

    def __init__(self, in_channels):
        super(InceptionE, self).__init__()
        self.branch1x1 = BasicConv2d(in_channels, 320, kernel_size=1)
        self.branch3x3_1 = BasicConv2d(in_channels, 384, kernel_size=1)
        self.branch3x3_2a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1))
        self.branch3x3_2b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0))
        self.branch3x3dbl_1 = BasicConv2d(in_channels, 448, kernel_size=1)
        self.branch3x3dbl_2 = BasicConv2d(448, 384, kernel_size=3, padding=1)
        self.branch3x3dbl_3a = BasicConv2d(384, 384, kernel_size=(1, 3), padding=(0, 1))
        self.branch3x3dbl_3b = BasicConv2d(384, 384, kernel_size=(3, 1), padding=(1, 0))
        self.branch_pool = BasicConv2d(in_channels, 192, kernel_size=1)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)
        branch3x3 = self.branch3x3_1(x)
        branch3x3 = [self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)]
        branch3x3 = torch.cat(branch3x3, 1)
        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = [self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)]
        branch3x3dbl = torch.cat(branch3x3dbl, 1)
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)
        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


class InceptionAux(nn.Module):

    def __init__(self, in_channels, num_classes):
        super(InceptionAux, self).__init__()
        self.conv0 = BasicConv2d(in_channels, 128, kernel_size=1)
        self.conv1 = BasicConv2d(128, 768, kernel_size=5)
        self.conv1.stddev = 0.01
        self.fc = nn.Linear(768, num_classes)
        self.fc.stddev = 0.001

    def forward(self, x):
        x = F.avg_pool2d(x, kernel_size=5, stride=3)
        x = self.conv0(x)
        x = self.conv1(x)
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


class conv_bn_relu(nn.Module):

    def __init__(self, in_channel, out_channel, stride=1, kszie=3, pad=0, has_bn=True, has_relu=True, bias=True, groups=1):
        """
        Basic block with one conv, one bn, one relu in series.

        Arguments
        ---------
        in_channel: int
            number of input channels
        out_channel: int
            number of output channels
        stride: int
            stride number
        kszie: int
            kernel size
        pad: int
            padding on each edge
        has_bn: bool
            use bn or not
        has_relu: bool
            use relu or not
        bias: bool
            conv has bias or not
        groups: int or str
            number of groups. To be forwarded to torch.nn.Conv2d
        """
        super(conv_bn_relu, self).__init__()
        self.conv = nn.Conv2d(in_channel, out_channel, kernel_size=kszie, stride=stride, padding=pad, bias=bias, groups=groups)
        if has_bn:
            self.bn = nn.BatchNorm2d(out_channel)
        else:
            self.bn = None
        if has_relu:
            self.relu = nn.ReLU()
        else:
            self.relu = None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class creat_residual_block(nn.Module):

    def __init__(self, inplanes, outplanes, stride, has_proj=False):
        super(creat_residual_block, self).__init__()
        self.has_proj = has_proj
        if self.has_proj:
            self.proj_conv = conv_bn_relu(inplanes, outplanes, stride=stride, kszie=1, pad=0, has_bn=True, has_relu=False, bias=False)
        self.conv1 = conv_bn_relu(inplanes, outplanes, stride=stride, kszie=3, pad=1, has_bn=True, has_relu=True, bias=False)
        self.conv2 = conv_bn_relu(outplanes, outplanes, stride=1, kszie=3, pad=1, has_bn=True, has_relu=False, bias=False)
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x
        if self.has_proj:
            residual = self.proj_conv(residual)
        x = self.conv1(x)
        x = self.conv2(x)
        x = x + residual
        x = self.relu(x)
        return x


class create_bottleneck(nn.Module):
    """
    Modified Bottleneck : We change the kernel size of projection conv from 1 to 3.

    """

    def __init__(self, inplanes, outplanes, stride, has_proj=False):
        super(create_bottleneck, self).__init__()
        self.has_proj = has_proj
        if self.has_proj:
            self.proj_conv = conv_bn_relu(inplanes, outplanes, stride=stride, kszie=3, pad=1, has_bn=True, has_relu=False, bias=False)
        self.conv1 = conv_bn_relu(inplanes, outplanes, stride=stride, kszie=3, pad=1, has_bn=True, has_relu=True, bias=False)
        self.conv2 = conv_bn_relu(outplanes, outplanes, stride=1, kszie=3, pad=1, has_bn=True, has_relu=True, bias=False)
        self.conv3 = conv_bn_relu(outplanes, outplanes, stride=1, kszie=3, pad=1, has_bn=True, has_relu=False, bias=False)
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x
        if self.has_proj:
            residual = self.proj_conv(residual)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x + residual
        x = self.relu(x)
        return x


class ShuffleNetV2(nn.Module):

    def __init__(self, stages_repeats, stages_out_channels, num_classes=1000, inverted_residual=InvertedResidual, crop_pad=4, fused_channls=[116, 232, 464], head_width=256, **kwargs):
        super(ShuffleNetV2, self).__init__()
        self.crop_pad = crop_pad
        if len(stages_repeats) != 3:
            raise ValueError('expected stages_repeats as list of 3 positive ints')
        if len(stages_out_channels) != 5:
            raise ValueError('expected stages_out_channels as list of 5 positive ints')
        self._stage_out_channels = stages_out_channels
        input_channels = 3
        output_channels = self._stage_out_channels[0]
        self.conv1 = nn.Sequential(nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False), nn.BatchNorm2d(output_channels), nn.ReLU(inplace=True))
        input_channels = output_channels
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]
        for name, repeats, output_channels in zip(stage_names, stages_repeats, self._stage_out_channels[1:]):
            false_stride = name != 'stage2'
            seq = [inverted_residual(input_channels, output_channels, 2, false_stride=false_stride)]
            for i in range(repeats - 1):
                seq.append(inverted_residual(output_channels, output_channels, 1))
            setattr(self, name, nn.Sequential(*seq))
            input_channels = output_channels
        output_channels = self._stage_out_channels[-1]
        self.channel_reduce = nn.Sequential(nn.Conv2d(sum(fused_channls), head_width, 1, 1, 0, bias=False), nn.BatchNorm2d(head_width), nn.ReLU(inplace=True))

    def _forward_impl(self, x):
        x = x / 255
        x_ch0 = (torch.unsqueeze(x[:, 2], 1) - 0.485) / 0.229
        x_ch1 = (torch.unsqueeze(x[:, 1], 1) - 0.456) / 0.224
        x_ch2 = (torch.unsqueeze(x[:, 0], 1) - 0.406) / 0.225
        x = torch.cat((x_ch0, x_ch1, x_ch2), 1)
        xs = []
        x = self.conv1(x)
        x = self.maxpool(x)
        x = self.stage2(x)
        xs.append(x)
        x = self.stage3(x)
        xs.append(x)
        x = self.stage4(x)
        xs.append(x)
        x = torch.cat(xs, 1)
        crop_pad = self.crop_pad
        x = x[:, :, crop_pad:-crop_pad, crop_pad:-crop_pad]
        x = self.channel_reduce(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


class upsample_block(nn.Module):
    """
    Upsample block. e.g. used for upsample and feature fusion in decoder
    """

    def __init__(self, h_channel, l_channel, out_channel):
        """
        h_channel:
            channel number of high-level feature

        l_channel:
            channel number of low-level feature

        out_channel:
            channel number of output feature after fusion
        """
        super(upsample_block, self).__init__()
        self.conv1 = conv_bn_relu(h_channel, out_channel, pad=1, bias=False)
        self.conv_adjust = conv_bn_relu(out_channel + l_channel, out_channel, pad=1, bias=False)

    def forward(self, high_level_f, low_level_f):
        """
        :param high_level_f: torch.Tensor
            high level feature with smaller resolution

        :param low_level_f: torch.Tensor
            low level feature with larger resolution

        Returns
        -------
        torch.Tensor
            feature fusion result
        """
        high_level_f = self.conv1(high_level_f)
        f_resize = F.interpolate(high_level_f, size=low_level_f.size()[2:], mode='bilinear', align_corners=False)
        f_fusion = torch.cat([f_resize, low_level_f], 1)
        f_adjust = self.conv_adjust(f_fusion)
        return f_adjust


class projector(nn.Module):
    """
    Projection layer to adjust channel number
    """

    def __init__(self, in_channel, out_channel):
        super(projector, self).__init__()
        self.conv1 = conv_bn_relu(in_channel, out_channel, pad=1, has_relu=False, bias=False)

    def forward(self, x):
        x = self.conv1(x)
        return x


def named_modules_with_dup(model: nn.Module, prefix: str='') ->Iterable[Tuple[str, nn.Module]]:
    """
    The same as `model.named_modules()`, except that it includes
    duplicated modules that have more than one name.
    """
    yield prefix, model
    for name, module in model._modules.items():
        if module is None:
            continue
        submodule_prefix = prefix + ('.' if prefix else '') + name
        yield from named_modules_with_dup(module, submodule_prefix)


def filter_reused_missing_keys(model: nn.Module, keys: List[str]) ->List[str]:
    """
    Filter "missing keys" to not include keys that have been loaded with another name.
    """
    keyset = set(keys)
    param_to_names = defaultdict(set)
    for module_prefix, module in named_modules_with_dup(model):
        for name, param in (list(module.named_parameters(recurse=False)) + list(module.named_buffers(recurse=False))):
            full_name = (module_prefix + '.' if module_prefix else '') + name
            param_to_names[param].add(full_name)
    for names in param_to_names.values():
        if any(n in keyset for n in names) and not all(n in keyset for n in names):
            [keyset.remove(n) for n in names if n in keyset]
    return list(keyset)


def group_checkpoint_keys(keys: List[str]) ->Dict[str, List[str]]:
    """
    Group keys based on common prefixes. A prefix is the string up to the final
    "." in each key.
    Args:
        keys (list[str]): list of parameter names, i.e. keys in the model
            checkpoint dict.
    Returns:
        dict[list]: keys with common prefixes are grouped into lists.
    """
    groups = defaultdict(list)
    for key in keys:
        pos = key.rfind('.')
        if pos >= 0:
            head, tail = key[:pos], [key[pos + 1:]]
        else:
            head, tail = key, []
        groups[head].extend(tail)
    return groups


def get_missing_parameters_message(keys: List[str]) ->str:
    """
    Get a logging-friendly message to report parameter names (keys) that are in
    the model but not found in a checkpoint.
    Args:
        keys (list[str]): List of keys that were not found in the checkpoint.
    Returns:
        str: message.
    """
    groups = group_checkpoint_keys(keys)
    msg = 'Some model parameters are not in the checkpoint:\n'
    for k, v in groups.items():
        msg += '{}:{}\n'.format(k, v)
    return msg


def get_unexpected_parameters_message(keys: List[str]) ->str:
    """
    Get a logging-friendly message to report parameter names (keys) that are in
    the checkpoint but not found in the model.
    Args:
        keys (list[str]): List of keys that were not found in the model.
    Returns:
        str: message.
    """
    groups = group_checkpoint_keys(keys)
    msg = 'The checkpoint contains parameters not used by the model:\n'
    for k, v in groups.items():
        msg += '{}:{}\n'.format(k, v)
    return msg


logger = logging.getLogger('global')


def md5sum(file_path) ->str:
    """Get md5sum string
    
    Parameters
    ----------
    file_path : str
        path to file to calculate md5sum
    
    Returns
    -------
    str
        md5 value string in hex
    """
    with open(file_path, 'rb') as f:
        md5sum_str = hashlib.md5(f.read()).hexdigest()
    return md5sum_str


class ModuleBase(nn.Module):
    """
    Module/component base class
    """
    default_hyper_params = dict(pretrain_model_path='')

    def __init__(self):
        super(ModuleBase, self).__init__()
        self._hyper_params = deepcopy(self.default_hyper_params)

    def get_hps(self) ->dict():
        """
        Getter function for hyper-parameters

        Returns
        -------
        dict
            hyper-parameters
        """
        return self._hyper_params

    def set_hps(self, hps: dict()) ->None:
        """
        Set hyper-parameters

        Arguments
        ---------
        hps: dict
            dict of hyper-parameters, the keys must in self.__hyper_params__
        """
        for key in hps:
            if key not in self._hyper_params:
                raise KeyError
            self._hyper_params[key] = hps[key]

    def update_params(self):
        model_file = self._hyper_params.get('pretrain_model_path', '')
        if model_file != '':
            state_dict = torch.load(model_file, map_location=torch.device('cpu'))
            if 'model_state_dict' in state_dict:
                state_dict = state_dict['model_state_dict']
            self.load_model_param(state_dict)
            logger.info('Load pretrained {} parameters from: {} whose md5sum is {}'.format(self.__class__.__name__, model_file, md5sum(model_file)))

    def load_model_param(self, checkpoint_state_dict):
        model_state_dict = self.state_dict()
        for k in list(checkpoint_state_dict.keys()):
            if k in model_state_dict:
                shape_model = tuple(model_state_dict[k].shape)
                shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
                if shape_model != shape_checkpoint:
                    logger.warning("'{}' has shape {} in the checkpoint but {} in the model! Skipped.".format(k, shape_checkpoint, shape_model))
                    checkpoint_state_dict.pop(k)
        incompatible = self.load_state_dict(checkpoint_state_dict, strict=False)
        if incompatible.missing_keys:
            missing_keys = filter_reused_missing_keys(self, incompatible.missing_keys)
            if missing_keys:
                logger.warning(get_missing_parameters_message(missing_keys))
        if incompatible.unexpected_keys:
            logger.warning(get_unexpected_parameters_message(incompatible.unexpected_keys))


class Registry(dict):
    """
    A helper class for managing registering modules, it extends a dictionary
    and provides a register functions.

    usually declared in XXX_base.py, e.g. videoanalyst/model/backbone/backbone_base.py

    used as decorator when declaring the module:

    @some_registry.register
    def foo():
        ...

    Access of module is just like using a dictionary, eg:
        f = some_registry["foo_module"]
    """

    def __init__(self, *args, **kwargs):
        self.name = 'Registry'
        if len(args) > 0 and isinstance(args[0], str):
            name, *args = args
            self.name = name
        super(Registry, self).__init__(*args, **kwargs)

    def register(self, module):
        name = module.__name__
        _register_generic(self, name, module)
        return module


TRACK_HEADS = Registry('TRACK_HEADS')


VOS_HEADS = Registry('VOS_HEADS')


def get_box(xy_ctr, offsets):
    offsets = offsets.permute(0, 2, 3, 1)
    offsets = offsets.reshape(offsets.shape[0], -1, 4)
    xy0 = xy_ctr[:, :, :] - offsets[:, :, :2]
    xy1 = xy_ctr[:, :, :] + offsets[:, :, 2:]
    bboxes_pred = torch.cat([xy0, xy1], 2)
    return bboxes_pred


def get_xy_ctr(score_size, score_offset, total_stride):
    batch, fm_height, fm_width = 1, score_size, score_size
    y_list = torch.linspace(0.0, fm_height - 1.0, fm_height).reshape(1, fm_height, 1, 1).repeat(1, 1, fm_width, 1)
    x_list = torch.linspace(0.0, fm_width - 1.0, fm_width).reshape(1, 1, fm_width, 1).repeat(1, fm_height, 1, 1)
    xy_list = score_offset + torch.cat([x_list, y_list], 3) * total_stride
    xy_ctr = xy_list.repeat(batch, 1, 1, 1).reshape(batch, -1, 2)
    xy_ctr = xy_ctr.type(torch.Tensor)
    return xy_ctr


class DenseboxHead(ModuleBase):
    """
    Densebox Head for siamfcpp

    Hyper-parameter
    ---------------
    total_stride: int
        stride in backbone
    score_size: int
        final feature map
    x_size: int
        search image size
    num_conv3x3: int
        number of conv3x3 tiled in head
    head_conv_bn: list
        has_bn flag of conv3x3 in head, list with length of num_conv3x3
    head_width: int
        feature width in head structure
    conv_weight_std: float
        std for conv init
    """
    default_hyper_params = dict(total_stride=8, score_size=17, x_size=303, num_conv3x3=3, head_conv_bn=[False, False, True], head_width=256, conv_weight_std=0.0001)

    def __init__(self):
        super(DenseboxHead, self).__init__()
        self.bi = torch.nn.Parameter(torch.tensor(0.0).type(torch.Tensor))
        self.si = torch.nn.Parameter(torch.tensor(1.0).type(torch.Tensor))
        self.cls_convs = []
        self.bbox_convs = []

    def forward(self, c_out, r_out):
        num_conv3x3 = self._hyper_params['num_conv3x3']
        cls = c_out
        bbox = r_out
        for i in range(0, num_conv3x3):
            cls = getattr(self, 'cls_p5_conv%d' % (i + 1))(cls)
            bbox = getattr(self, 'bbox_p5_conv%d' % (i + 1))(bbox)
        cls_score = self.cls_score_p5(cls)
        cls_score = cls_score.permute(0, 2, 3, 1)
        cls_score = cls_score.reshape(cls_score.shape[0], -1, 1)
        ctr_score = self.ctr_score_p5(cls)
        ctr_score = ctr_score.permute(0, 2, 3, 1)
        ctr_score = ctr_score.reshape(ctr_score.shape[0], -1, 1)
        offsets = self.bbox_offsets_p5(bbox)
        offsets = torch.exp(self.si * offsets + self.bi) * self.total_stride
        self.fm_ctr = self.fm_ctr
        bbox = get_box(self.fm_ctr, offsets)
        return [cls_score, ctr_score, bbox, cls]

    def update_params(self):
        x_size = self._hyper_params['x_size']
        score_size = self._hyper_params['score_size']
        total_stride = self._hyper_params['total_stride']
        score_offset = (x_size - 1 - (score_size - 1) * total_stride) // 2
        self._hyper_params['score_offset'] = score_offset
        self.score_size = self._hyper_params['score_size']
        self.total_stride = self._hyper_params['total_stride']
        self.score_offset = self._hyper_params['score_offset']
        ctr = get_xy_ctr(self.score_size, self.score_offset, self.total_stride)
        self.fm_ctr = ctr
        self.fm_ctr.require_grad = False
        self._make_conv3x3()
        self._make_conv_output()
        self._initialize_conv()

    def _make_conv3x3(self):
        num_conv3x3 = self._hyper_params['num_conv3x3']
        head_conv_bn = self._hyper_params['head_conv_bn']
        head_width = self._hyper_params['head_width']
        self.cls_conv3x3_list = []
        self.bbox_conv3x3_list = []
        for i in range(num_conv3x3):
            cls_conv3x3 = conv_bn_relu(head_width, head_width, stride=1, kszie=3, pad=0, has_bn=head_conv_bn[i])
            bbox_conv3x3 = conv_bn_relu(head_width, head_width, stride=1, kszie=3, pad=0, has_bn=head_conv_bn[i])
            setattr(self, 'cls_p5_conv%d' % (i + 1), cls_conv3x3)
            setattr(self, 'bbox_p5_conv%d' % (i + 1), bbox_conv3x3)
            self.cls_conv3x3_list.append(cls_conv3x3)
            self.bbox_conv3x3_list.append(bbox_conv3x3)

    def _make_conv_output(self):
        head_width = self._hyper_params['head_width']
        self.cls_score_p5 = conv_bn_relu(head_width, 1, stride=1, kszie=1, pad=0, has_relu=False)
        self.ctr_score_p5 = conv_bn_relu(head_width, 1, stride=1, kszie=1, pad=0, has_relu=False)
        self.bbox_offsets_p5 = conv_bn_relu(head_width, 4, stride=1, kszie=1, pad=0, has_relu=False)

    def _initialize_conv(self):
        num_conv3x3 = self._hyper_params['num_conv3x3']
        conv_weight_std = self._hyper_params['conv_weight_std']
        conv_list = []
        for i in range(num_conv3x3):
            conv_list.append(getattr(self, 'cls_p5_conv%d' % (i + 1)).conv)
            conv_list.append(getattr(self, 'bbox_p5_conv%d' % (i + 1)).conv)
        conv_list.append(self.cls_score_p5.conv)
        conv_list.append(self.ctr_score_p5.conv)
        conv_list.append(self.bbox_offsets_p5.conv)
        conv_classifier = [self.cls_score_p5.conv]
        assert all(elem in conv_list for elem in conv_classifier)
        num_classes = 1
        pi = 0.01
        bv = -np.log((1 - pi) / pi)
        for ith in range(len(conv_list)):
            conv = conv_list[ith]
            torch.nn.init.normal_(conv.weight, std=conv_weight_std)
            if conv in conv_classifier:
                torch.nn.init.constant_(conv.bias, torch.tensor(bv))
            else:
                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(conv.weight)
                bound = 1 / np.sqrt(fan_in)
                nn.init.uniform_(conv.bias, -bound, bound)


class DecoderHead(ModuleBase):
    """
    DecoderHead for SAT

    Hyper-parameter
    ---------------
    output_size: int
        output size of predicted mask

    """
    default_hyper_params = dict(output_size=257, input_channel_list=[512, 256, 128, 64])

    def __init__(self):
        super(DecoderHead, self).__init__()
        self.output_size = self._hyper_params['output_size']
        self.out_projector = projector(128, 1)
        self.f_s16_projector = projector(256, 1)
        self.f_s8_projector = projector(256, 1)
        self.activation = nn.Sigmoid()

    def update_params(self):
        input_channel_list = self._hyper_params['input_channel_list']
        self.upblock1 = upsample_block(input_channel_list[0], input_channel_list[0], 256)
        self.upblock2 = upsample_block(256, input_channel_list[1], 256)
        self.upblock3 = upsample_block(256, input_channel_list[2], 256)
        self.upblock4 = upsample_block(256, input_channel_list[3], 128)

    def forward(self, feature_list, phase='train'):
        x1, x2, x3, x4, x5 = feature_list
        f_s32 = self.upblock1(x1, x2)
        f_s16 = self.upblock2(f_s32, x3)
        f_s8 = self.upblock3(f_s16, x4)
        f_s4 = self.upblock4(f_s8, x5)
        p = self.out_projector(f_s4)
        p_resize = F.interpolate(p, (self.output_size, self.output_size), mode='bilinear', align_corners=False)
        if phase == 'train':
            pred_s16 = self.f_s16_projector(f_s16)
            pred_s16_resize = F.interpolate(pred_s16, (self.output_size, self.output_size), mode='bilinear', align_corners=False)
            pred_s8 = self.f_s8_projector(f_s8)
            pred_s8_resize = F.interpolate(pred_s8, (self.output_size, self.output_size), mode='bilinear', align_corners=False)
            return [pred_s16_resize, pred_s8_resize, p_resize]
        else:
            prediction = self.activation(p_resize)
            return prediction


VOS_TASKMODELS = Registry('VOS_TASKMODELS')


class SatVOS(ModuleBase):
    """
    State-Aware Tracker model for VOS

    Hyper-Parameters
    ----------------
    pretrain_model_path: string
        path to parameter to be loaded into module
    """
    default_hyper_params = dict(pretrain_model_path='')

    def __init__(self, GML_extractor, joint_encoder, decoder, loss):
        super(SatVOS, self).__init__()
        self.GML_extractor = GML_extractor
        self.joint_encoder = joint_encoder
        self.decoder = decoder
        self.loss = loss

    def forward(self, *args, phase='train'):
        """
        Perform VOS process for different phases (e.g. train / global_feature / segment)

        Arguments
        ---------
        filterd_image: torch.Tensor
            filtered image patch for global modeling loop

        saliency_image: torch.Tensor
            saliency image for saliency encoder
        corr_feature: torch.Tensor
            correlated feature produced by siamese encoder
        global_feature: torch.Tensor
            global feature produced by global modeling loop

        Returns
        -------
        f_g: torch.Tensor
            global feature extracted from filtered image
        pred_mask: torch.Tensor
            predicted mask after sigmoid for the patch of saliency image

        """
        if phase == 'train':
            saliency_image, corr_feature, filtered_image = args
            global_feature = self.GML_extractor(filtered_image)
            enc_features = self.joint_encoder(saliency_image, corr_feature)
            decoder_features = [global_feature] + enc_features
            out_list = self.decoder(decoder_features, phase='train')
        elif phase == 'global_feature':
            filtered_image, = args
            f_g = self.GML_extractor(filtered_image)
            out_list = [f_g]
        elif phase == 'segment':
            saliency_image, corr_feature, global_feature = args
            enc_features = self.joint_encoder(saliency_image, corr_feature)
            decoder_features = [global_feature] + enc_features
            outputs = self.decoder(decoder_features, phase='test')
            pred_mask = outputs
            out_list = [pred_mask]
        else:
            raise ValueError('Phase non-implemented.')
        return out_list

    def set_device(self, dev):
        if not isinstance(dev, torch.device):
            dev = torch.device(dev)
        self
        if self.loss is not None:
            for loss_name in self.loss:
                self.loss[loss_name]


TRACK_TASKMODELS = Registry('TRACK_TASKMODELS')


class SiamTrack(ModuleBase):
    """
    SiamTrack model for tracking

    Hyper-Parameters
    ----------------
    pretrain_model_path: string
        path to parameter to be loaded into module
    head_width: int
        feature width in head structure
    """
    default_hyper_params = dict(pretrain_model_path='', head_width=256, conv_weight_std=0.01, neck_conv_bias=[True, True, True, True], corr_fea_output=False)

    def __init__(self, backbone, head, loss=None):
        super(SiamTrack, self).__init__()
        self.basemodel = backbone
        self.head = head
        self.loss = loss

    def forward(self, *args, phase='train'):
        """
        Perform tracking process for different phases (e.g. train / init / track)

        Arguments
        ---------
        target_img: torch.Tensor
            target template image patch
        search_img: torch.Tensor
            search region image patch

        Returns
        -------
        fcos_score_final: torch.Tensor
            predicted score for bboxes, shape=(B, HW, 1)
        fcos_bbox_final: torch.Tensor
            predicted bbox in the crop, shape=(B, HW, 4)
        fcos_cls_prob_final: torch.Tensor
            classification score, shape=(B, HW, 1)
        fcos_ctr_prob_final: torch.Tensor
            center-ness score, shape=(B, HW, 1)
        """
        if phase == 'train':
            training_data = args[0]
            target_img = training_data['im_z']
            search_img = training_data['im_x']
            f_z = self.basemodel(target_img)
            f_x = self.basemodel(search_img)
            c_z_k = self.c_z_k(f_z)
            r_z_k = self.r_z_k(f_z)
            c_x = self.c_x(f_x)
            r_x = self.r_x(f_x)
            r_out = xcorr_depthwise(r_x, r_z_k)
            c_out = xcorr_depthwise(c_x, c_z_k)
            fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(c_out, r_out)
            predict_data = dict(cls_pred=fcos_cls_score_final, ctr_pred=fcos_ctr_score_final, box_pred=fcos_bbox_final)
            if self._hyper_params['corr_fea_output']:
                predict_data['corr_fea'] = corr_fea
            return predict_data
        elif phase == 'feature':
            target_img, = args
            f_z = self.basemodel(target_img)
            c_z_k = self.c_z_k(f_z)
            r_z_k = self.r_z_k(f_z)
            out_list = [c_z_k, r_z_k]
        elif phase == 'track':
            if len(args) == 3:
                search_img, c_z_k, r_z_k = args
                f_x = self.basemodel(search_img)
                c_x = self.c_x(f_x)
                r_x = self.r_x(f_x)
            elif len(args) == 4:
                c_z_k, r_z_k, c_x, r_x = args
            else:
                raise ValueError('Illegal args length: %d' % len(args))
            r_out = xcorr_depthwise(r_x, r_z_k)
            c_out = xcorr_depthwise(c_x, c_z_k)
            fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(c_out, r_out)
            fcos_cls_prob_final = torch.sigmoid(fcos_cls_score_final)
            fcos_ctr_prob_final = torch.sigmoid(fcos_ctr_score_final)
            fcos_score_final = fcos_cls_prob_final * fcos_ctr_prob_final
            extra = dict(c_x=c_x, r_x=r_x, corr_fea=corr_fea)
            out_list = fcos_score_final, fcos_bbox_final, fcos_cls_prob_final, fcos_ctr_prob_final, extra
        else:
            raise ValueError('Phase non-implemented.')
        return out_list

    def update_params(self):
        """
        Load model parameters
        """
        self._make_convs()
        self._initialize_conv()
        super().update_params()

    def _make_convs(self):
        head_width = self._hyper_params['head_width']
        self.r_z_k = conv_bn_relu(head_width, head_width, 1, 3, 0, has_relu=False)
        self.c_z_k = conv_bn_relu(head_width, head_width, 1, 3, 0, has_relu=False)
        self.r_x = conv_bn_relu(head_width, head_width, 1, 3, 0, has_relu=False)
        self.c_x = conv_bn_relu(head_width, head_width, 1, 3, 0, has_relu=False)

    def _initialize_conv(self):
        conv_weight_std = self._hyper_params['conv_weight_std']
        conv_list = [self.r_z_k.conv, self.c_z_k.conv, self.r_x.conv, self.c_x.conv]
        for ith in range(len(conv_list)):
            conv = conv_list[ith]
            torch.nn.init.normal_(conv.weight, std=conv_weight_std)

    def set_device(self, dev):
        if not isinstance(dev, torch.device):
            dev = torch.device(dev)
        self
        if self.loss is not None:
            for loss_name in self.loss:
                self.loss[loss_name]


class ContrastiveLoss(torch.nn.Module):
    """
    Contrastive loss function.
    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
    """

    def __init__(self, margin=2.0):
        super(ContrastiveLoss, self).__init__()
        self.margin = margin

    def forward(self, output1, output2, label):
        euclidean_distance = F.pairwise_distance(output1, output2)
        loss_contrastive = torch.mean((1 - label) * torch.pow(euclidean_distance, 2) + label * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
        return loss_contrastive


class SiameseNetwork(nn.Module):

    def __init__(self):
        super(SiameseNetwork, self).__init__()
        self.cnn1 = nn.Sequential(nn.ReflectionPad2d(1), nn.Conv2d(1, 4, kernel_size=3), nn.ReLU(inplace=True), nn.BatchNorm2d(4), nn.Dropout2d(p=0.2), nn.ReflectionPad2d(1), nn.Conv2d(4, 8, kernel_size=3), nn.ReLU(inplace=True), nn.BatchNorm2d(8), nn.Dropout2d(p=0.2), nn.ReflectionPad2d(1), nn.Conv2d(8, 8, kernel_size=3), nn.ReLU(inplace=True), nn.BatchNorm2d(8), nn.Dropout2d(p=0.2))
        self.fc1 = nn.Sequential(nn.Linear(8 * 100 * 100, 500), nn.ReLU(inplace=True), nn.Linear(500, 500), nn.ReLU(inplace=True), nn.Linear(500, 5))

    def forward_once(self, x):
        output = self.cnn1(x)
        output = output.view(output.size()[0], -1)
        output = self.fc1(output)
        return output

    def forward(self, input1, input2):
        output1 = self.forward_once(input1)
        output2 = self.forward_once(input2)
        return output1, output2


class ConvBlock(nn.Module):

    def __init__(self, inplane, outplane, kernel_size, padding, stride):
        super(ConvBlock, self).__init__()
        self.convmodule = nn.Sequential(nn.Conv2d(inplane, outplane, kernel_size=kernel_size, padding=padding, stride=stride, bias=False), nn.BatchNorm2d(outplane), nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.convmodule(x)
        return x


class Mask(nn.Module):

    def __init__(self):
        super(Mask, self).__init__()

    def forward(self, z_f, x_f):
        raise NotImplementedError

    def template(self, template):
        raise NotImplementedError

    def track(self, search):
        raise NotImplementedError

    def param_groups(self, start_lr, feature_mult=1):
        params = filter(lambda x: x.requires_grad, self.parameters())
        params = [{'params': params, 'lr': start_lr * feature_mult}]
        return params


class MaskCorr(DepthwiseXCorr):

    def __init__(self, in_channels, hidden, out_channels, kernel_size=3, hidden_kernel_size=5):
        super(MaskCorr, self).__init__(in_channels, hidden, out_channels, kernel_size, hidden_kernel_size)

    def forward(self, kernel, search):
        kernel = self.conv_kernel(kernel)
        search = self.conv_search(search)
        feature = xcorr_depthwise(search, kernel)
        out = self.head(feature)
        return out, feature


class MultiMaskCorr(Mask):

    def __init__(self, oSz=63, weighted=False):
        super(MultiMaskCorr, self).__init__()
        self.weighted = weighted
        for i in range(3):
            self.add_module('maskcorr' + str(i + 2), MaskCorr())
        if self.weighted:
            self.mask_weight = nn.Parameter(torch.ones(3))

    def forward(self, z_fs, x_fs):
        mask = []
        for idx, (z_f, x_f) in enumerate(zip(z_fs, x_fs), start=2):
            maskcorr = getattr(self, 'maskcorr' + str(idx))
            m = maskcorr(z_f, x_f)
            mask.append(m)
        if self.weighted:
            mask_weight = F.softmax(self.mask_weight, 0)

        def avg(lst):
            return sum(lst) / len(lst)

        def weighted_avg(lst, weight):
            s = 0
            for i in range(len(weight)):
                s += lst[i] * weight[i]
            return s
        if self.weighted:
            return weighted_avg(mask, mask_weight)
        else:
            return avg(mask)

    def forward_corrs(self, z_fs, x_fs):
        mask = []
        for idx, (z_f, x_f) in enumerate(zip(z_fs, x_fs), start=2):
            maskcorr = getattr(self, 'maskcorr' + str(idx))
            m = maskcorr.mask.forward_corr(z_f, x_f)
            mask.append(m)
        if self.weighted:
            mask_weight = F.softmax(self.mask_weight, 0)

        def avg(lst):
            return sum(lst) / len(lst)

        def weighted_avg(lst, weight):
            s = 0
            for i in range(len(weight)):
                s += lst[i] * weight[i]
            return s
        if self.weighted:
            return weighted_avg(mask, mask_weight)
        else:
            return avg(mask)


class Refine(nn.Module):

    def __init__(self):
        super(Refine, self).__init__()
        self.v0 = nn.Sequential(nn.Conv2d(64, 16, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(16, 4, 3, padding=1), nn.ReLU(inplace=True))
        self.v1 = nn.Sequential(nn.Conv2d(256, 64, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(64, 16, 3, padding=1), nn.ReLU(inplace=True))
        self.v2 = nn.Sequential(nn.Conv2d(512, 128, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(128, 32, 3, padding=1), nn.ReLU(inplace=True))
        self.h2 = nn.Sequential(nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(inplace=True))
        self.h1 = nn.Sequential(nn.Conv2d(16, 16, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(16, 16, 3, padding=1), nn.ReLU(inplace=True))
        self.h0 = nn.Sequential(nn.Conv2d(4, 4, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(4, 4, 3, padding=1), nn.ReLU(inplace=True))
        self.deconv = nn.ConvTranspose2d(256, 32, 15, 15)
        self.post0 = nn.Conv2d(32, 16, 3, padding=1)
        self.post1 = nn.Conv2d(16, 4, 3, padding=1)
        self.post2 = nn.Conv2d(4, 1, 3, padding=1)

    def forward(self, f, corr_feature, pos):
        p0 = F.pad(f[0], [16, 16, 16, 16])[:, :, 4 * pos[0]:4 * pos[0] + 61, 4 * pos[1]:4 * pos[1] + 61]
        p1 = F.pad(f[1], [8, 8, 8, 8])[:, :, 2 * pos[0]:2 * pos[0] + 31, 2 * pos[1]:2 * pos[1] + 31]
        p2 = F.pad(f[2], [4, 4, 4, 4])[:, :, pos[0]:pos[0] + 15, pos[1]:pos[1] + 15]
        p3 = corr_feature[:, :, pos[0], pos[1]].view(-1, 256, 1, 1)
        out = self.deconv(p3)
        out = self.post0(F.upsample(self.h2(out) + self.v2(p2), size=(31, 31)))
        out = self.post1(F.upsample(self.h1(out) + self.v1(p1), size=(61, 61)))
        out = self.post2(F.upsample(self.h0(out) + self.v0(p0), size=(127, 127)))
        out = out.view(-1, 127 * 127)
        return out


class RPN(nn.Module):

    def __init__(self):
        super(RPN, self).__init__()

    def forward(self, z_f, x_f):
        raise NotImplementedError


class UPChannelRPN(RPN):

    def __init__(self, anchor_num=5, in_channels=256, out_channels=0):
        super(UPChannelRPN, self).__init__()
        cls_output = 2 * anchor_num
        loc_output = 4 * anchor_num
        self.template_cls_conv = nn.Conv2d(in_channels, in_channels * cls_output, kernel_size=3)
        self.template_loc_conv = nn.Conv2d(in_channels, in_channels * loc_output, kernel_size=3)
        self.search_cls_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3)
        self.search_loc_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3)
        self.loc_adjust = nn.Conv2d(loc_output, loc_output, kernel_size=1)

    def forward(self, z_f, x_f):
        cls_kernel = self.template_cls_conv(z_f)
        loc_kernel = self.template_loc_conv(z_f)
        cls_feature = self.search_cls_conv(x_f)
        loc_feature = self.search_loc_conv(x_f)
        cls = xcorr_fast(cls_feature, cls_kernel)
        loc = self.loc_adjust(xcorr_fast(loc_feature, loc_kernel))
        return cls, loc


class DepthwiseRPN(RPN):

    def __init__(self, anchor_num=5, in_channels=256, out_channels=256):
        super(DepthwiseRPN, self).__init__()
        self.cls = DepthwiseXCorr(in_channels, out_channels, 2 * anchor_num)
        self.loc = DepthwiseXCorr(in_channels, out_channels, 4 * anchor_num)

    def forward(self, z_f, x_f):
        cls = self.cls(z_f, x_f)
        loc = self.loc(z_f, x_f)
        return cls, loc


class MultiRPN(RPN):

    def __init__(self, anchor_num, in_channels, weighted=False):
        super(MultiRPN, self).__init__()
        self.weighted = weighted
        for i in range(len(in_channels)):
            self.add_module('rpn' + str(i + 2), DepthwiseRPN(anchor_num, in_channels[i], in_channels[i]))
        if self.weighted:
            self.cls_weight = nn.Parameter(torch.ones(len(in_channels)))
            self.loc_weight = nn.Parameter(torch.ones(len(in_channels)))

    def forward(self, z_fs, x_fs):
        cls = []
        loc = []
        for idx, (z_f, x_f) in enumerate(zip(z_fs, x_fs), start=2):
            rpn = getattr(self, 'rpn' + str(idx))
            c, l = rpn(z_f, x_f)
            cls.append(c)
            loc.append(l)
        if self.weighted:
            cls_weight = F.softmax(self.cls_weight, 0)
            loc_weight = F.softmax(self.loc_weight, 0)

        def avg(lst):
            return sum(lst) / len(lst)

        def weighted_avg(lst, weight):
            s = 0
            for i in range(len(weight)):
                s += lst[i] * weight[i]
            return s
        if self.weighted:
            return weighted_avg(cls, cls_weight), weighted_avg(loc, loc_weight)
        else:
            return avg(cls), avg(loc)


class PositionEmbeddingLearned(nn.Module):
    """
    Absolute pos embedding, learned.
    """

    def __init__(self, num_pos_feats=256):
        super().__init__()
        self.row_embed = nn.Embedding(50, num_pos_feats)
        self.col_embed = nn.Embedding(50, num_pos_feats)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.uniform_(self.row_embed.weight)
        nn.init.uniform_(self.col_embed.weight)

    def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors
        h, w = x.shape[-2:]
        i = torch.arange(w, device=x.device)
        j = torch.arange(h, device=x.device)
        x_emb = self.col_embed(i)
        y_emb = self.row_embed(j)
        pos = torch.cat([x_emb.unsqueeze(0).repeat(h, 1, 1), y_emb.unsqueeze(1).repeat(1, w, 1)], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
        return pos


class UpdateResNet256(nn.Module):

    def __init__(self, config=None):
        super(UpdateResNet256, self).__init__()
        self.update = nn.Sequential(nn.Conv2d(768, 96, 1), nn.ReLU(inplace=True), nn.Conv2d(96, 256, 1))

    def forward(self, x, x0):
        response = self.update(x)
        response += x0
        return response


class UpdateResNet512(nn.Module):

    def __init__(self, config=None):
        super(UpdateResNet512, self).__init__()
        self.update = nn.Sequential(nn.Conv2d(1536, 192, 1), nn.ReLU(inplace=True), nn.Conv2d(192, 512, 1))

    def forward(self, x, x0):
        response = self.update(x)
        response += x0
        return response


class UpdateNet(nn.Module):

    def __init__(self, config=None):
        super(UpdateNet, self).__init__()
        self.update = nn.Sequential(nn.Conv2d(768, 96, 1), nn.ReLU(inplace=True), nn.Conv2d(96, 256, 1))

    def forward(self, x):
        response = self.update(x)
        return response


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (AdjustAllLayer,
     lambda: ([], {'in_channels': [4, 4], 'out_channels': [4, 4]}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (AdjustLayer,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (AlexNet,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (AlexNetLegacy,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (AlexNetV0,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (AlexNetV1,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (BalancedLoss,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicBlock,
     lambda: ([], {'inplanes': 4, 'planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicBlock_C,
     lambda: ([], {'in_planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicConv2d,
     lambda: ([], {'in_channels': 4, 'out_channels': 4, 'kernel_size': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicConv2d_1x1,
     lambda: ([], {'in_channels': 4, 'out_channels': 4, 'kernel_size': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicConv2d_3x3,
     lambda: ([], {'inplanes': 4, 'planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Bottleneck,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 96, 64, 64])], {}),
     True),
    (CAModule,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 64, 4, 4])], {}),
     True),
    (ContrastiveLoss,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (ConvBlock,
     lambda: ([], {'inplane': 4, 'outplane': 4, 'kernel_size': 4, 'padding': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (DepthwiseRPN,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 256, 64, 64]), torch.rand([4, 256, 64, 64])], {}),
     True),
    (DepthwiseXCorr,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (FocalLoss,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (GHMCLoss,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (Incep22,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (InceptionA,
     lambda: ([], {'in_channels': 4, 'pool_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InceptionAux,
     lambda: ([], {'in_channels': 4, 'num_classes': 4}),
     lambda: ([torch.rand([4, 4, 64, 64])], {}),
     True),
    (InceptionB,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InceptionC,
     lambda: ([], {'in_channels': 4, 'channels_7x7': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InceptionD,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InceptionE,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InceptionM,
     lambda: ([], {'in_channels': 4, 'planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (InvertedResidual,
     lambda: ([], {'inp': 4, 'oup': 4, 'stride': 1, 'expand_ratio': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MobileNetV2,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (MobileNetV3,
     lambda: ([], {'cfgs': _mock_config(), 'mode': 'large'}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     True),
    (PositionEmbeddingSine,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (ResNeXt,
     lambda: ([], {'num_blocks': [4, 4], 'cardinality': 4, 'bottleneck_width': 4}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (ResNeXt22,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 128, 128])], {}),
     True),
    (SELayer,
     lambda: ([], {'channel': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (SiamFC,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (TransformerDecoderLayer,
     lambda: ([], {'d_model': 4, 'nhead': 4}),
     lambda: ([torch.rand([4, 4]), torch.rand([4, 4])], {}),
     True),
    (TransformerEncoderLayer,
     lambda: ([], {'d_model': 4, 'nhead': 4}),
     lambda: ([torch.rand([4, 4])], {}),
     True),
    (UPChannelBAN,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 256, 64, 64]), torch.rand([4, 256, 64, 64])], {}),
     True),
    (UPChannelRPN,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 256, 64, 64]), torch.rand([4, 256, 64, 64])], {}),
     True),
    (UpdateNet,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 768, 64, 64])], {}),
     True),
    (UpdateResNet256,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 768, 64, 64]), torch.rand([4, 256, 64, 64])], {}),
     True),
    (UpdateResNet512,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 1536, 64, 64]), torch.rand([4, 512, 64, 64])], {}),
     True),
    (_BatchNorm2d,
     lambda: ([], {'num_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (conv_bn_relu,
     lambda: ([], {'in_channel': 4, 'out_channel': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (creat_residual_block,
     lambda: ([], {'inplanes': 4, 'outplanes': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (create_bottleneck,
     lambda: ([], {'inplanes': 4, 'outplanes': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (h_sigmoid,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (h_swish,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (projector,
     lambda: ([], {'in_channel': 4, 'out_channel': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (upsample_block,
     lambda: ([], {'h_channel': 4, 'l_channel': 4, 'out_channel': 4}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     True),
]

class Test_HonglinChu_SiamTrackers(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

    def test_004(self):
        self._check(*TESTCASES[4])

    def test_005(self):
        self._check(*TESTCASES[5])

    def test_006(self):
        self._check(*TESTCASES[6])

    def test_007(self):
        self._check(*TESTCASES[7])

    def test_008(self):
        self._check(*TESTCASES[8])

    def test_009(self):
        self._check(*TESTCASES[9])

    def test_010(self):
        self._check(*TESTCASES[10])

    def test_011(self):
        self._check(*TESTCASES[11])

    def test_012(self):
        self._check(*TESTCASES[12])

    def test_013(self):
        self._check(*TESTCASES[13])

    def test_014(self):
        self._check(*TESTCASES[14])

    def test_015(self):
        self._check(*TESTCASES[15])

    def test_016(self):
        self._check(*TESTCASES[16])

    def test_017(self):
        self._check(*TESTCASES[17])

    def test_018(self):
        self._check(*TESTCASES[18])

    def test_019(self):
        self._check(*TESTCASES[19])

    def test_020(self):
        self._check(*TESTCASES[20])

    def test_021(self):
        self._check(*TESTCASES[21])

    def test_022(self):
        self._check(*TESTCASES[22])

    def test_023(self):
        self._check(*TESTCASES[23])

    def test_024(self):
        self._check(*TESTCASES[24])

    def test_025(self):
        self._check(*TESTCASES[25])

    def test_026(self):
        self._check(*TESTCASES[26])

    def test_027(self):
        self._check(*TESTCASES[27])

    def test_028(self):
        self._check(*TESTCASES[28])

    def test_029(self):
        self._check(*TESTCASES[29])

    def test_030(self):
        self._check(*TESTCASES[30])

    def test_031(self):
        self._check(*TESTCASES[31])

    def test_032(self):
        self._check(*TESTCASES[32])

    def test_033(self):
        self._check(*TESTCASES[33])

    def test_034(self):
        self._check(*TESTCASES[34])

    def test_035(self):
        self._check(*TESTCASES[35])

    def test_036(self):
        self._check(*TESTCASES[36])

    def test_037(self):
        self._check(*TESTCASES[37])

    def test_038(self):
        self._check(*TESTCASES[38])

    def test_039(self):
        self._check(*TESTCASES[39])

    def test_040(self):
        self._check(*TESTCASES[40])

    def test_041(self):
        self._check(*TESTCASES[41])

    def test_042(self):
        self._check(*TESTCASES[42])

    def test_043(self):
        self._check(*TESTCASES[43])

    def test_044(self):
        self._check(*TESTCASES[44])

    def test_045(self):
        self._check(*TESTCASES[45])

    def test_046(self):
        self._check(*TESTCASES[46])

    def test_047(self):
        self._check(*TESTCASES[47])

    def test_048(self):
        self._check(*TESTCASES[48])

    def test_049(self):
        self._check(*TESTCASES[49])

    def test_050(self):
        self._check(*TESTCASES[50])

