【LUT技术专题】ECLUT代码解读

目录

原文概要

1. 训练

2. 转表

3. 测试

---

本文是对ECLUT技术的代码解读，原文解读请看ECLUT。 

原文概要

ECLUT通过EC模块增大网络感受野，提升超分效果，实现SRLUT的改进，主要是2个创新点：

• 提出了一个扩展卷积（expand Conv），跟RCLUT一样，是一个用于低成本提升RF大小的模块。
• 一个简单有效的scaling方法来减小旋转累计的量化误差，该方法的计算成本基本可以忽略不计。

其网络结构图如下：

​

Expanded Conv：

​

首先根据前面的讲解，我们可以发现ECLUT是在SRLUT的基础上进行的优化，ECLUT因为只使用了2个查询点，因此LUT的size在全尺寸下也不会很大。作者是采用了1的间隔来实现ECLUT，所以精度是能够得到保证的，流程就可以分为，训练、转表以及推理，忽略确保精度的微调步骤。

---

1. 训练

代码实现如下：

### A lightweight deep network ###
class SRNet(torch.nn.Module):def __init__(self, upscale=4):super(SRNet, self).__init__()self.upscale = upscaleself.conv1 = nn.Conv2d(1, 64, [2,1], stride=1, padding=0, dilation=1)self.conv2 = nn.Conv2d(64, 64, 1, stride=1, padding=0, dilation=1)self.conv3 = nn.Conv2d(64, 64, 1, stride=1, padding=0, dilation=1)self.conv4 = nn.Conv2d(64, 64, 1, stride=1, padding=0, dilation=1)self.conv5 = nn.Conv2d(64, 64, 1, stride=1, padding=0, dilation=1)self.conv6 = nn.Conv2d(64, 9*upscale*upscale, 1, stride=1, padding=0, dilation=1)self.pixel_shuffle = nn.PixelShuffle(upscale)self.expand_conv = ExpandConv()# Init weightsfor m in self.modules():classname = m.__class__.__name__if classname.lower().find('conv') != -1:nn.init.kaiming_normal(m.weight)nn.init.constant(m.bias, 0)elif classname.find('bn') != -1:m.weight.data.normal_(1.0, 0.02)m.bias.data.fill_(0)def forward(self, x_in, alpha):B, C, H, W = x_in.size()x_in = x_in.reshape(B*C, 1, H, W)x = self.conv1(x_in)x = self.conv2(F.relu(x))x = self.conv3(F.relu(x))x = self.conv4(F.relu(x))x = self.conv5(F.relu(x))x = self.conv6(F.relu(x))x = self.pixel_shuffle(x)x = self.expand_conv(x)x = x.reshape(B, C, self.upscale*(H-1), self.upscale*W)return x * alpha

根据前面的讲解，ECLUT是在SRLUT的基础上进行的优化，因此整体结构的实现博主采取了与SRLUT的代码差不多的方式。不同之处在于：

1）ECLUT因为只使用了2个查询点，因此首层的卷积kernel_size是2x1。

2）pixel_shuffle之前的通道多了9倍，因为需要进行我们讲到的EC模块，即expand_conv。

3）最后的alpha加权，减小量化误差。

这里关于最重要的模块expand_conv的实现如下：

class ExpandConv(torch.nn.Module):def __init__(self):super(ExpandConv, self).__init__()passdef forward(self, inputs):# b, 9, h*scale, w*scale# (0,0) -> (2, 2)# pad left right top bottomshift_x1 = F.pad(inputs[:, 0:1, 1:, 1:], (0,1,0,1), mode='reflect')shift_x2 = F.pad(inputs[:, 1:2, 1:, :], (0,0,0,1), mode='reflect')shift_x3 = F.pad(inputs[:, 2:3, 1:, :-1], (1,0,0,1), mode='reflect')shift_x4 = F.pad(inputs[:, 3:4, :, 1:], (0,1,0,0), mode='reflect')shift_x5 = inputs[:, 4:5, :, :]shift_x6 = F.pad(inputs[:, 5:6, :, :-1], (1,0,0,0), mode='reflect')shift_x7 = F.pad(inputs[:, 6:7, :-1, 1:], (0,1,1,0), mode='reflect')shift_x8 = F.pad(inputs[:, 7:8, :-1, :], (0,0,1,0), mode='reflect')shift_x9 = F.pad(inputs[:, 8:9, :-1, :-1], (1,0,1,0), mode='reflect')result = shift_x1 + shift_x2 + shift_x3 +\shift_x4 + shift_x5 + shift_x6 + shift_x7 + shift_x8 + shift_x9result /= 9return result

对inputs进行9个方向的偏移，然后pad对应的长度使得其输出大小不变，最后将他们相加进行一个融合，完成了这个EC模块的计算过程。

---

2. 转表

首先我们需要初始化一个2D的输入。

def get_ec_lut(interval):# 1D inputbase = torch.arange(0, 257, interval)  # 0-256base[-1] -= 1L = base.size(0)# 2D input# 256*256   0 0 0...    |1 1 1...     |...|255 255 255...first = base.cuda().unsqueeze(1).repeat(1, L).reshape(-1)# 256*256   0 1 2 .. 255|0 1 2 ... 255|...|0 1 2 ... 255second = base.cuda().repeat(L)onebytwo = torch.stack([first, second], 1)  # [256*256, 2]# Rearange input: [N, 4] -> [N, C=1, H=2, W=2]input_tensor = onebytwo.unsqueeze(1).unsqueeze(1).reshape(-1, 1, 2, 1).float() / 255.0return input_tensor

然后送入推理即可，因为ECLUT是通过对多输出进行偏移得到更大的RF，因此这里就已经得到了最终所需要的LUT。

---

3. 测试

测试的部分，针对第一个卷积及其后面的1x1卷积部分，可以转换为一个完全精度的2D LUT，这个查询是不需要插值的，因此十分简单，第二个是expand_conv，只需要我们将1.训练中实现的torch版本改写为numpy的版本即可，对于最后的alpha加权，使用加2再右移2位来实现。这个代码的实现在ECLUT的讲解中有提到。

---

以上针对于ECLUT代码实现的部分讲解完毕，如果有不清楚的问题欢迎大家提出。