手写实现LRN局部响应归一化算子

1、重写算子的需求

芯片推理过程中遇到很多算子计算结果不对的情况，原因是封装的算子会在某些特殊情况下计算超限，比如输入shape特别大或者数值特别大时，LRN算子计算会出现NAN值，所以需要重写算子。先对输入数据做一个预处理，计算后再在合适的地方转换回去。

2、lrn算子的原理

LRN全称是local response normalization，局部响应归一化，想了解原理的点这个AlexNet原论文。

官方API伪代码如下：

sqr_sum[a, b, c, d] = sum(input[a,b, c, d - depth_radius : d + depth_radius + 1] ** 2)
output = input / (bias +alpha * sqr_sum) ** beta
1
2

在alexnet的原论文中，输入为 [batch_size, 224, 224, 96]，这里224×224是图片的大小，经过第一次卷积再经过ReLU，就是LRN函数的输入。

注意上面API说明里的sum函数，意思就是，可能解释起来比较拗口，针对batch里每一个图的后3维向量，[224, 224, d - depth_radius : d + depth_radius + 1]，对它按照半径 depth_radius求每个图里的每个像素的平方，再把这2× depth_radius+1个平方过后的图片加起来，就得到了这个batch的sqr_sum。

3、手写实现lrn算子

下面参考原论文和pytorch源码，实现自己手写的lrn算子，其中avg_pool3d就是实现了按照半径 depth_radius求每个图里的每个像素的平方：

def custom_lrn(input_tensor, N=5, alpha=1e-4, beta=0.75):
    x_sq = torch.square(input_tensor).unsqueeze(1)
    sizes = input_tensor.size()
    x_reshape = x_sq.view(sizes[0], 1, sizes[1], sizes[2], -1)
    x_pad = torch.nn.functional.pad(x_reshape, (0,0,0,0,2,2))
    x_pool3d = torch.nn.functional.avg_pool3d(x_pad, (N, 1, 1),stride=1).squeeze(1)
    x_squeeze = x_pool3d.view(sizes)
    x_scale = torch.mul(x_squeeze, alpha) + (1.0)
    x_scale_pow = torch.pow(x_scale, beta)
    out = input_tensor / x_scale_pow
    return out
1
2
3
4
5
6
7
8
9
10
11

测试一下和pytorch实现的官方API的结果情况：

import torch
import numpy as np
inputs = torch.randn(1, 64, 56, 56, dtype=torch.float32)*20
SEED = 1
def set_seed(seed=1):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
set_seed(SEED)

def custom_lrn(input_tensor, N=5, alpha=1e-4, beta=0.75):
    x_sq = torch.square(input_tensor).unsqueeze(1)
    sizes = input_tensor.size()
    x_reshape = x_sq.view(sizes[0], 1, sizes[1], sizes[2], -1)
    x_pad = torch.nn.functional.pad(x_reshape, (0,0,0,0,2,2))
    x_pool3d = torch.nn.functional.avg_pool3d(x_pad, (N, 1, 1),stride=1).squeeze(1)
    x_squeeze = x_pool3d.view(sizes)
    x_scale = torch.mul(x_squeeze, alpha) + (1.0)
    x_scale_pow = torch.pow(x_scale, beta)
    out = input_tensor / x_scale_pow
    return out

lrn2 = torch.nn.functional.local_response_norm(inputs, size=5)


# print(custom_lrn(inputs))
# print(lrn2)
print('custom_lrn与pytorch官方的lrn算子是否相等：',torch.allclose(custom_lrn(inputs), lrn2))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

测试输出结果完全一致，说明此算子与官方实现的算子是一致的。

>> custom_lrn与pytorch官方的lrn算子是否相等：True
1