# 一般张量要设置了梯度需求才会在后续 backward 中保存梯度
x = x.requires_grad_(True)

# 分类的模型和损失函数，在计算期间产生的张量如 y 和 loss 都会保存一个 grad_fn，指向用于计算当前层反向传播的 Function 对象
y = model.forward(x)
loss = target_loss(y)

# 执行 backward() 会沿着 grad_fn 逐步往前计算梯度
# 注意，如果没有设置 retain_grad() 的话会在 backward() 之后清空计算图的前向缓存，也就是这次涉及的相关节点都不能再进行反向传播了
loss.backward()

# 对梯度后处理，转换成 opencv 图像用来观察 loss 的产生原因对应到输入的区域
saliency_map = torch.sum(torch.abs(x.grad), axis=1, keepdims=True)
saliency_map = saliency_map.detach().cpu().numpy()
saliency_map = cvglue.scale_min_max(saliency_map, 0, 255, percentile=99)
saliency_map = np.uint8(saliency_map[0].transpose((1,2,0)))

nn.Sequential(*[Conv(), BN(), ReLU()])

torch.nn.RNN(input_size, hidden_size, num_layers=1, batch_first=False)

a = torch.ones(5)
b = a.numpy()   # a 和 b 共用一块内存
a = np.ones(5)
b = torch.from_numpy(a)  # 不要用 torch.Tensor(a)，速度会非常慢

def affine2theta(warp_mat, in_size, out_size):
    W1 = in_size[3]
    H1 = in_size[2]
    W2 = out_size[3]
    H2 = out_size[2]
    warp_wh_2 = np.array([2/W2, 0, -1, 0, 2/H2, -1, 0, 0, 1]).reshape(3,3)
    warp_wh_1 = np.array([W1/2, 0, W1/2, 0, H1/2, H1/2, 0, 0, 1]).reshape(3,3)
    warp_mat_r = np.insert(warp_mat, [2], [0,0,1], axis=0)
    aff_theta = torch.from_numpy(np.linalg.inv(warp_wh_2 @ warp_mat_r @ warp_wh_1))[:2].unsqueeze(0).type(torch.float32)
    return aff_theta 

# src (1xCxHxW)  out_size (4)
W1 = src.size()[3]
H1 = src.size()[2]
W2 = out_size[3]
H2 = out_size[2]
# keep resolution when shinking
if W1 > W2 and H1 > H2:
    out_size = src.size()
# aff_theta (1x2x3)
aff_theta = affine2theta(M, src.size(), out_size)
# grid (1xHxWx2)
grid = torch.nn.functional.affine_grid(aff_theta.cuda(), out_size)
dst = torch.nn.functional.grid_sample(src.cuda(), grid)
dst = dst[:,:,:H2,:W2]

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        # nn.Module子类的函数必须在构造函数中执行父类的构造函数
        super(Net, self).__init__()
        
        # 卷积层 '1'表示输入图片为单通道, '6'表示输出通道数，'3'表示卷积核为3*3
        self.conv1 = nn.Conv2d(1, 6, 3) 
        #线性层，输入1350个特征，输出10个特征
        self.fc1   = nn.Linear(1350, 10)  #这里的1350是如何计算的呢？这就要看后面的forward函数
    #正向传播 
    def forward(self, x): 
        print(x.size()) # 结果：[1, 1, 32, 32]
        # 卷积 -> 激活 -> 池化 
        x = self.conv1(x) #根据卷积的尺寸计算公式，计算结果是30，具体计算公式后面第二章第四节 卷积神经网络 有详细介绍。
        x = F.relu(x)
        print(x.size()) # 结果：[1, 6, 30, 30]
        x = F.max_pool2d(x, (2, 2)) #我们使用池化层，计算结果是15
        x = F.relu(x)
        print(x.size()) # 结果：[1, 6, 15, 15]
        # reshape，‘-1’表示自适应
        #这里做的就是压扁的操作 就是把后面的[1, 6, 15, 15]压扁，变为 [1, 1350]
        x = x.view(x.size()[0], -1) 
        print(x.size()) # 这里就是fc1层的的输入1350 
        x = self.fc1(x)        
        return x

net = Net()
print(net)

for name,parameters in net.named_parameters():
    print(name,':',parameters.size())

torch.cuda.synchronize()   # 等待之前的设备上计算完成
start = timer()

y = model.forward(x)

torch.cuda.synchronize()   # 等待模型在设备上的计算完成
end = timer()

class SmoothOneHot(nn.Module):
    def __init__(self, classes, target=None, smooth_val=None):
        super().__init__()
        if target is None:
            target = torch.tensor(1.5)
        self.target = nn.Parameter(data=target)

        if smooth_val is None:
            smooth_val = torch.ones(1)
        self.smooth_val = nn.Parameter(data=smooth_val)
        
    def forward(self, target_center):
        res = torch.exp(-(target_center - self.target)**2 / (2 * self.smooth_val**2))
        smooth_one_hot = res / res.sum(dim=-1, keepdim=True)
        return smooth_one_hot

def init_seeds(seed=0):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if seed == 0:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

from torch.utils.data import Dataset
import pandas as pd
# 定义一个数据集
class BulldozerDataset(Dataset):
    """ 数据集演示 """
    def __init__(self, csv_file):
        """实现初始化方法，在初始化的时候将数据读载入"""
        self.df=pd.read_csv(csv_file)
    def __len__(self):
        '''
        返回df的长度，决定数据集的大小，若返回小于 len(self.df) 的值则会截断数据集
        '''
        return len(self.df)
    def __getitem__(self, idx):
        '''
        根据 idx 返回一行数据
        '''
        return self.df.iloc[idx].SalePrice
		# 也可以返回字典，dataloader 会自动将里面的元素 batch 起来
		# return {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, 'feat': feat_tensor, 'path': A_path}
        
ds_demo= BulldozerDataset('median_benchmark.csv')
# 实现了 __len__ 方法所以可以直接使用len获取数据总数
len(ds_demo)
# 用索引可以直接访问对应的数据, 对应 __getitem__ 方法
ds_demo[0]

# 返回迭代器
dl = torch.utils.data.DataLoader(ds_demo, batch_size=10, shuffle=True, num_workers=0)

# 通过迭代枚举来获得数据
for i, data in enumerate(dl):
	# 通过给 enumerate 设置 start=n 可以指定下标 i 的起始位置，但数据 dl 依然从头开始遍历
    print(i, data)

from tfrecord.torch.dataset import TFRecordDataset
index_path = None
description = {"depth": "int", "image_raw": "byte"}
# 返回 Dataset
dataset = TFRecordDataset(tfrecord_path, index_path, description) 
# 用 dataloader 封装
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1)
# 用 PIL decode 图像文件，转换为 Tensor
for i, data in enumerate(dataloader):
    x = Image.open(io.BytesIO(data['image_raw'].numpy()))
    y = transforms.Compose([transforms.ToTensor()])(x)

# 也可以用 albumentations，功能更强大，但多少有些坑
import torchvision.transforms as transforms
transform_A = transforms.Compose([transforms.ToTensor()])

# 1. 用 PIL 预处理，最后转成 Tensor
for i, data in enumerate(dataloader):
    x = Image.open(io.BytesIO(data['image_raw'].numpy()))
    y = transform_A(x)

# 2. 用 OpenCV 预处理，最后转成 Tensor
for i, data in enumerate(dataloader):
    x = cv2.imdecode(data['image_raw'].numpy(), cv2.IMREAD_COLOR)
    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
    y = transform_A(x)
    
# 3. 用 numpy 预处理，最后转成 Tensor
torch.from_array(a)

import torch.optim as optim

# create your optimizer
optimizer = optim.SGD(net.parameters(), lr=0.01)
	
# in your training loop:
# for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()   # zero the gradient buffers, 不然会累积
output = net(input)
loss = criterion(output, target)
loss.backward()  # 也可以直接网络输出 output 为 loss，然后 output.backward()
optimizer.step()    # Does the update

# 使用方法非常简单，只要套起来就行
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)

# 数据分发，并行计算
y = model(x)

# 调用模型的方法
y = model.module.xxx(x)

import random
import nunmpy as np
# 可以配合 amp 进行混合精度训练，进一步加快速度，更多资料见参考 [5]、[6]
from torch.cuda.amp import autocast

# 首先要设置当前进程使用的 CUDA 设备，以及初始化分布式环境，设置了之后
# 统一使用 .cuda() 就能将 Tensor 或模型等放到相应设备上
torch.cuda.set_device(opt.local_rank)
torch.distributed.init_process_group(
    'nccl',
    init_method='env://'
)

# 设定当前进程的全局 rank，用于指定一个进程输出日志和保存模型
opt.rank = torch.distributed.get_rank()

# 如果用到其他库的随机函数的话，在 fork 多进程模式下需要重新设置其种子，
# 不然会出现不同进程产生相同随机数列的情况
random.seed(opt.rank)
np.random.seed(seed=opt.rank)

# 为了每个 worker 使用不同的数据来训练，需要 DistributedSampler 来进行
# 划分（官方提供的该函数是非连续的划分），它会使 enumerate(dataloader) 
# 得到的 index 范围变为 [0, ITERS/WORLD_SIZE]，也就是每个 worker 被分配
# 到其中一部分数据来进行训练。也可以用自己写的连续划分的 sampler
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)

# 如果需要将数据处理也放到 GPU 的话可以考虑 DALI 之类的库
dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=opt.batchSize,
            drop_last=True,	   # 会扔掉 (batchsize * 卡数) 无法整除的余数，扔掉的通过
			                   # shuffle 可以取到，具体看 len(dataset) 实现是否包含
            num_workers=int(opt.nThreads),
            sampler=train_sampler,
            shuffle=not opt.serial_batches and (train_sampler is None))


# （可选）将模型里的 BN 转成 SyncBN，参数会 copy，这里为了抵消数据方差减小可以适当增大 lr
# from apex.parallel import convert_syncbn_model
# model = convert_syncbn_model(model)

# 分布式包装，指定本地显卡 local_rank
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[opt.local_rank], output_device=opt.local_rank)

# amp 的缩放器初始化
scaler = torch.cuda.amp.GradScaler()

# 恢复训练需要注意的是 torch.load 时要将读进来的数据通过 map_location
# 映射到当前 worker 的机器上，因为模型一般只在某台机器上保存，读取时
# 默认会加载到保存时所用的机器上
if opt.continue_train:
    checkpoint_ = torch.load(os.path.join(save_dir, 'latest_checkpoint.pt'), map_location=lambda storage, loc: storage.cuda(opt.local_rank))
    optimizer_G.load_state_dict(checkpoint_['optimizer_G'])
    optimizer_D.load_state_dict(checkpoint_['optimizer_D'])
    scaler.load_state_dict(checkpoint_['amp'])  # 读取 loss scale 和 unskipped 数

# 同步各机器的 loss 的函数，worker 在这里会锁住等待其他 worker 运行到这里同步
# 得到均值返回，注意不要只放到 rank==0 上做，不然其他 worker 进不来会卡死 0
def reduce_tensor(tensor):
    rt = tensor.clone()
    torch.distributed.all_reduce(rt, op=torch.distributed.ReduceOp.SUM)
    rt /= torch.distributed.get_world_size()
    return rt

for epoch in range(num_epochs):
	train_sampler.set_epoch(epoch)   # 设置 epoch 作为 shuffle 的种子，
                                     # 使每 epoch 不会读到相同顺序的数据，
                                     # 但似乎会让不同实验取间数据的顺序完全相同

    for i, data in enumerate(dataloader):# 各卡经过 sampler 得到相同 batchsize 数据，
	                                    # 每次遍历只抽取 dataset 中的 (bs * 卡数) 整除
										# 若数据集大小=10, droplast, bs=3，卡数=2，
										# 则每 epoch 取 6 个数据，len(dataloader)=1
    	with autocast(enabled=opt.fp16):
			loss = model(data)
        
    	optimizer_G.zero_grad()
		scaler.scale(loss_G).backward()
		scaler.step(optimizer_G)
        
        # 注意计算图是否有重用，有则需要 retain_graph
        optimizer_D.zero_grad()
        scaler.scale(loss_D).backward()
		scaler.step(optimizer_D)
		
		scaler.update()
	    
    	# 同步 loss 值用于显示输出
		loss_i = reduce_tensor(loss.data).item()
        
        # 打印输出，保存 checkpoint 等只在 rank==0 上做
        if opt.rank == 0:
        	if (i + print_freq_delta) % opt.print_freq == 0:
        		print(loss_i)
                
    if opt.rank == 0:
        if epoch % opt.save_epoch_freq == 0:
            checkpoint_amp = {
                'optimizer_G': optimizer_G.state_dict(),
                'optimizer_D': optimizer_D.state_dict(),
                'amp': scaler.state_dict()   # 记录当前 loss scale 和 unskipped 数
            }
            torch.save(checkpoint_amp, os.path.join(save_dir, 'latest_checkpoint.pt'))
            model.module.save(epoch, save_latest=True)

CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 train.py

# 先初始化
hvd.init()

# 和其他方法一样，一般让线程和 GPU 一一对应
torch.cuda.set_device(hvd.local_rank())

# 用分布式优化器封装，将会通过 allreduce 或者 allgather 将梯度均值化
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_params())

# 还要从 rank 0 广播初始状态到其他进程
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

# 旧版需要指定名字作为 reduce 的索引 key，新版默认用生成的自增名字
def reduce_tensor(tensor, name):
    rt = tensor.clone()
	avg_rt = hvd.allreduce(rt, name=name)
    return avg_rt

# 先取得需要修改的模块
for module in model.modules():
    if isinstance(module, nn.Conv2d):
        list_conv += [module]
    elif isinstance(module, nn.BatchNorm2d):
        list_bn += [module]
    if len(list_conv) == 3:
        break

# 对于卷积需要修改输入输出通道和具体的 weights 和 bias（如果有），最后重新初始化参数
list_conv[0].in_channels = in_channels
list_conv[0].out_channels = out_channels
weight0 = list_conv[0].weight.detach()
weight0 = torch.Tensor(
    list_conv[0].out_channels,
    list_conv[0].in_channels // list_conv[0].groups,
    *list_conv[0].kernel_size
)
list_conv[0].weight = nn.parameter.Parameter(weight0)
list_conv[0].reset_parameters()  # 参数初始化

# 对于 BN，需要修改 weights 和 bias，同时还要修改 num_features 和 running mean/var
# 因为 running mean/var 视 BN 参数可能为 parameter 也可能是 buffer，所以需要按照具体情况而修改
weight3 = list_bn[0].weight.detach()
bias3 = list_bn[0].bias.detach()
weight3 = torch.Tensor(out_channels)
bias3 = torch.Tensor(out_channels)
list_bn[0].num_features = out_channels
list_bn[0].weight = nn.parameter.Parameter(weight3)
list_bn[0].bias = nn.parameter.Parameter(bias3)
list_bn[0].register_buffer('running_mean', torch.zeros(out_channels))
list_bn[0].register_buffer('running_var', torch.ones(out_channels))
list_bn[0].register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
list_bn[0].reset_parameters()     # 参数初始化

# 读取参数
pretrained_dict = pretrained.state_dict() 
model_dict = m.state_dict() 
 
# 找寻 pretrained_dict 中与 model_dict 相同名称的模块参数
update_dict =  {k: v for k, v in pretrained_dict.items() if k in model_dict} 
 
# 使用这些参数更新 model_dict
model_dict.update(update_dict) 
 
# 将 model_dict 写回自定义模型
m.load_state_dict(model_dict)

from torch.optim.lr_scheduler import _LRScheduler
class WarmupMultiStepLR(_LRScheduler):
    def __init__(self, optimizer, warmup_factor=1.0 / 3, warmup_iters=500, last_epoch=-1):
        self.warmup_factor = warmup_factor
        self.warmup_iters = warmup_iters
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        warmup_factor = 1
        if self.last_epoch < self.warmup_iters:
            alpha = float(self.last_epoch) / self.warmup_iters
            warmup_factor = self.warmup_factor * (1 - alpha) + alpha
        return [
            base_lr
            * warmup_factor
            for base_lr in self.base_lrs
        ]

scheduler = WarmupMultiStepLR(optimizer)

for IN TRAINING LOOP:
    ***
    optimizer.step()
    ***
    scheduler.step()
    ***

def model(nn.Module):
   xxx
   
model.eval()    # 设置为前向模式，对于 BatchNorm 以及 Dropout 等有用的设置
                # 类似的，训练时需要 model.train() 设为训练模式
               
model.cuda()    # 使用 GPU，同 model.to("cuda") 和 DataParallel

# DataParallel 也会把模型放到 GPU 上
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)

# 不要用 torch.Tensor(y).unsqueeze(0).cuda() 来转换输入 Tensor，贼慢
trans = Alb.Compose([ToTensorV2(transpose_mask=True)])
# normalization 之类批量运算的也可以在 GPU 上进行，加快速度
x = trans(image=img)['image'].unsqueeze(0).cuda()

# 需要指明不保留梯度，节省显存
with torch.no_grad():
    # 输入数据也要在同一个 device 上
	y = model.forward(x)
    
# 使用 GPU 后处理提升速度
y = (y + 1)*127.5
# 不要用 y.type(dtype=torch.ByteTensor) 直接转换到 cpu，会很慢
y_array = y.type(dtype=torch.cuda.ByteTensor)

# 拿到 cpu 数据
y_images = y_array.cpu().numpy()

torch.cuda.set_device(0)
x = torch.rand(1,3,3).cuda()
torch.save(x, 'xxx.pth')   # cuda:0 保存

# 读取到 cuda:0
a = torch.load('xxx.pth')

# 读取到 cuda:1
b = torch.load('xxx.pth', map_location=lambda storage, loc: storage.cuda(1))

# 读取到 cpu
c = torch.load('xxx.pth', map_location=lambda storage, loc: storage)

from thop import profile
from thop import clever_format
# 自定义的层需要手写统计代码
model = ResNet18()
inputx = torch.randn(1, 1, 28, 28)
flops, params = profile(model, inputs=(inputx, ))
flops, params = clever_format([flops, params], "%.3f")

# 或者使用 torchsummaryX
from torchsummaryX import summary
summary(model, torch.zeros((1, 3, 256, 256)))

print("moduel_name", "channel_id", "channel_sum", "largest_channel_sum", "module", sep='\t\t')
for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
        sum_k = torch.mean(torch.abs(module.weight), axis=[1,2,3])
        norm_sum_k = sum_k / sum_k.max()
        for v in range(norm_sum_k.shape[0]):
            if torch.abs(norm_sum_k[v]) < 1e-3:
                print(name, v, "%.4f"%float(sum_k[v]), "%.4f"%float(sum_k.max()), module, sep='\t\t')
    elif isinstance(module, nn.BatchNorm2d):
        sum_k = torch.abs(module.weight)
        norm_sum_k = sum_k / sum_k.max()
        for v in range(norm_sum_k.shape[0]):
            if torch.abs(norm_sum_k[v]) < 1e-3:
                print(name, v, "%.4f"%float(sum_k[v]), "%.4f"%float(sum_k.max()), module, sep='\t\t')

# 未测试
def prune(model, amount=0.3):  
    # Prune model to requested global sparsity  
    import torch.nn.utils.prune as prune  
    print('Pruning model... ', end='')  
    for name, m in model.named_modules():  
        if isinstance(m, nn.Conv2d):  
            prune.l1_unstructured(m, name='weight', amount=amount)  # prune  
            prune.remove(m, 'weight')  # make permanent  
    print(' %.3g global sparsity' % sparsity(model))

dummy_input = torch.randn(1, 1, 28, 28, device='cuda')
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, dummy_input, "LeNet_MNIST.onnx", verbose=True, input_names=input_names, output_names=output_names)

x = torch.rand((1, 3, 1024, 576)).cuda()
model_trt_1 = torch2trt(model, [x], fp16_mode=True, max_batch_size=8)
# model_trt_1 = torch2trt(model, [x], int8_mode=True, int8_calib_dataset=calid_dataset, max_batch_size=1, int8_calib_batch_size=8)
torch.save(model_trt_1.state_dict(), 'localL_extra123vertblurre_T4trt_1024x576_fp16.pth')

import torch.backends.cudnn as cudnn

torch.set_grad_enabled(False)
cudnn.benchmark = True

import cv2
import copy
import numpy as np


def crop_rect(img, rect, border_mode=cv2.BORDER_CONSTANT, value=0):
    """Crop a rectangle area from image, this allow the area partly outside the image.
    Args:
        img(array):           opencv image
        rect(tuple/array):    tuple of (left, top, right, bottom)
    Outs:
        face_img(array)
        pad_l, pad_t(tuple):  bias of axis origin
    """
    if isinstance(rect, np.ndarray):
        rect_int = np.int64(rect.flatten())
    else:
        rect_int = list(map(int, rect))
    face_lt = rect_int[:2]
    face_rb = rect_int[2:4]
    pad_l = -face_lt[0] if face_lt[0] < 0 else 0
    pad_t = -face_lt[1] if face_lt[1] < 0 else 0
    pad_r = face_rb[0] - img.shape[1] if face_rb[0] > img.shape[1] else 0
    pad_b = face_rb[1] - img.shape[0] if face_rb[1] > img.shape[0] else 0
    pad_img = cv2.copyMakeBorder(
        img, pad_t, pad_b, pad_l, pad_r, border_mode, value=value
    )
    face_img = pad_img[
        face_lt[1] + pad_t : face_rb[1] + pad_t + pad_b,
        face_lt[0] + pad_l : face_rb[0] + pad_l + pad_r,
        :,
    ]
    return face_img, (pad_l, pad_t)


def gradual_mean(raw_array, nan_mask=None, max_window_size=32):
    """
    raw_array:   (H, W, C) array
    nan_mask:    (H, W) array
    """
    assert len(raw_array.shape) == 3
    src_array = copy.deepcopy(raw_array)
    if nan_mask is None:
        nan_mask = np.sum(src_array, axis=-1)
    nan_ys, nan_xs = np.where(np.isnan(nan_mask))
    for cx, cy in zip(nan_xs, nan_ys):
        window_size = max_window_size
        window_mean = []
        while window_size > 1:
            half_l = window_size // 2
            window_rect = (cx - half_l, cy - half_l, cx + window_size, cy + window_size)
            window_img = crop_rect(src_array, window_rect)[0]
            window_mean.append(np.nanmean(window_img, axis=(0, 1)))
            if np.isnan(np.sum(window_mean[-1], axis=-1)):
                break
            window_size = window_size // 2
        window_mean = np.array(window_mean)
        src_array[cy, cx] = window_mean.mean(axis=0)
    return src_array

import copy
import warnings
import numpy as np
import torch
from torch.nn import functional as F


def gradual_mean_pytorch(raw_tensor, max_window_size=33):
    """
    Compute the gradual mean of a tensor using PyTorch functions.
    Args:
        raw_tensor:       a PyTorch tensor of shape (batch_size, channels, height, width)
                        containing the input data
        max_window_size:  an integer specifying the maximum size of the filter window
    Returns:
        raw_tensor_fil:    a PyTorch tensor of the same shape as raw_tensor
                            containing the filtered data
    """

    # 1. Calculate filters window size
    max_iters = int(np.log2(max_window_size - 1))
    window_sizes = np.logspace(1, max_iters, max_iters, base=2, dtype=np.int64) + 1
    if window_sizes[-1] != max_window_size:
        warnings.warn(
            f"max_window_size {max_window_size} is not aligned with max filter size {window_sizes[-1]}",
            UserWarning,
        )

    # 2. Get non-nan tensor and mask
    nan_mask = torch.isnan(raw_tensor)
    valid_mask = torch.logical_not(nan_mask)
    valid_tensor = torch.nan_to_num(raw_tensor, nan=0.0)

    # 3. Smoothen non-nan tensor and mar
    filtered_tensor = torch.cat(
        [
            F.avg_pool2d(valid_tensor, ws, stride=1, padding=ws // 2)
            / F.avg_pool2d(valid_mask.float(), ws, stride=1, padding=ws // 2)
            for ws in window_sizes
        ],
        dim=0,
    )

    # 4. In case there were still nan
    valid_filtered_mask = torch.logical_not(torch.isnan(filtered_tensor))
    valid_filtered_tensor = torch.nan_to_num(filtered_tensor, nan=0.0)

    # 5. Mean merge all filtered tensor
    mean_filtered_tensor = torch.sum(
        valid_filtered_tensor, dim=0, keepdims=True
    ) / torch.sum(valid_filtered_mask, dim=0, keepdims=True)

    # 6. Replace NaN values in the input tensor
    raw_tensor_filtered = copy.deepcopy(raw_tensor)
    raw_tensor_filtered[nan_mask] = mean_filtered_tensor[nan_mask]

    if torch.sum(torch.isnan(raw_tensor_filtered)) > 0:
        print("nan still exist after filtering.")

    return raw_tensor_filtered