本文是對4DLUT技術(shù)的代碼解讀，原文解讀請看4DLUT文章講解。

1、原文概要

4D LUT針對 3DLUT 在局部處理效果上的局限性，提出優(yōu)化方案。其核心亮點在于：通過引入圖像上下文編碼器提取上下文信息，并將該信息作為額外通道與原圖組成 4 通道輸入，從而使 4DLUT 能夠基于上下文感知實現(xiàn)圖像增強。
4DLUT的整體流程如下所示，可以分為4步：

1. 首先使用上下文編碼器（Context Encoder）通過端到端學習從輸入圖像生成表示像素級類別的上下文映射。
2. 同時，利用參數(shù)編碼器（Parameter Encoder）生成圖像自適應系數(shù)，用于融合可學習的預定義基礎4DLUTs（Basis 4D LUTs）。
3. 然后基于參數(shù)編碼器的輸出，使用4D LUTs融合模塊（4D LUTS Fusion）將可學習的基礎4DLUTs整合成最終具有更多增強功能的上下文感知4D LUT。
4. 最后，利用組合context map的RGBC通道圖像使用4DLUT進行插值得到增強的圖像。

2、代碼結(jié)構(gòu)

代碼整體結(jié)構(gòu)如下：

quadrilinear_cpp是四次插值的cpp代碼實現(xiàn)，datasets.py中放了數(shù)據(jù)集的處理，models_x.py放了模型和損失函數(shù)的實現(xiàn)，train.py放了整體訓練的流程。

3 、核心代碼模塊

models_x.py 文件

這個文件包含了4DLUT中context encoder、param encoder、四次插值的實現(xiàn)和兩個正則損失（平滑損失和單調(diào)損失）的計算。

1. Generator_for_info類

此為context encoder的實現(xiàn)。

class Generator_for_info(nn.Module):def __init__(self, in_channels=3):super(Generator_for_info, self).__init__()self.input_layer = nn.Sequential(nn.Conv2d(in_channels, 16, 3, stride=1, padding=1),nn.LeakyReLU(0.2),nn.InstanceNorm2d(16, affine=True),)self.mid_layer = nn.Sequential(*generator_block(16, 16, normalization=True),*generator_block(16, 16, normalization=True),*generator_block(16, 16, normalization=True),)self.output_layer = nn.Sequential(nn.Dropout(p=0.5),nn.Conv2d(16, 1, 3, stride=1, padding=1),nn.Sigmoid())def forward(self, img_input):x = self.input_layer(img_input)identity = xout = self.mid_layer(x)out += identityout = self.output_layer(out)return out

可以看到，輸入經(jīng)過一系列卷積和歸一化模塊，最終經(jīng)過一個輸出通道為1的卷積和一個sigmoid激活函數(shù)得到context map。

其中的discriminator_block實現(xiàn)如下：

def discriminator_block(in_filters, out_filters, normalization=False):"""Returns downsampling layers of each discriminator block"""layers = [nn.Conv2d(in_filters, out_filters, 3, stride=2, padding=1)]layers.append(nn.LeakyReLU(0.2))if normalization:layers.append(nn.InstanceNorm2d(out_filters, affine=True))#layers.append(nn.BatchNorm2d(out_filters))return layers

其實就是一個簡單的卷積，搭配了一個激活函數(shù)，根據(jù)normalization選項的不同插入InstanceNorm。

2. Generator_for_bias

該類實現(xiàn)了4DLUT中的param encoder類：

class Generator_for_bias(nn.Module):def __init__(self, in_channels=3):super(Generator_for_bias, self).__init__()self.model = nn.Sequential(nn.Upsample(size=(256,256),mode='bilinear'),nn.Conv2d(3, 16, 3, stride=2, padding=1),nn.LeakyReLU(0.2),nn.InstanceNorm2d(16, affine=True),*discriminator_block(16, 32, normalization=True),*discriminator_block(32, 64, normalization=True),*discriminator_block(64, 128, normalization=True),*discriminator_block(128, 128),nn.Dropout(p=0.5),nn.Conv2d(128, 12, 8, padding=0),)def forward(self, img_input):return self.model(img_input)

首先通過一個resize模塊將輸入降采樣至256分辨率，隨后經(jīng)過一系列卷積塊處理，最終通過一個kernel_size為8的12通道輸出卷積層，生成12個用于后續(xù)LUT融合的參數(shù)。

3. Generator4DLUT_identity

該類實現(xiàn)4DLUT的初始化。

class Generator4DLUT_identity(nn.Module):def __init__(self, dim=17):super(Generator4DLUT_identity, self).__init__()if dim == 17:file = open("Identity4DLUT17.txt", 'r')elif dim == 33:file = open("Identity4DLUT33.txt", 'r')lines = file.readlines()buffer = np.zeros((3,2,dim,dim,dim), dtype=np.float32)for p in range(0,2):for i in range(0,dim):for j in range(0,dim):for k in range(0,dim):n = p * dim*dim*dim + i * dim*dim + j*dim + kx = lines[n].split()buffer[0,p,i,j,k] = float(x[0])buffer[1,p,i,j,k] = float(x[1])buffer[2,p,i,j,k] = float(x[2])self.LUT_en = nn.Parameter(torch.from_numpy(buffer).requires_grad_(True))self.QuadrilinearInterpolation_4D = QuadrilinearInterpolation_4D()def forward(self, x):_, output = self.QuadrilinearInterpolation_4D(self.LUT_en, x)return output

這里的dim參數(shù)決定了LUT采樣時的bins數(shù)量，dim值越大采樣精度越高。博主存在一個疑問：context的維度參數(shù)p應當與dim保持一致，而非默認的2，否則context map只能獲取兩個離散值，后果是嚴重限制上下文信息的表達能力。實際操作中是將Identity4DLUT文本數(shù)據(jù)加載到buffer進行初始化，生成LUT_en用于后續(xù)的四維差值計算。

4. TV_4D

該類實現(xiàn)的是兩個正則化的損失函數(shù)。

class TV_4D(nn.Module):def __init__(self, dim=17):super(TV_4D,self).__init__()self.weight_r = torch.ones(3,2,dim,dim,dim-1, dtype=torch.float)self.weight_r[:,:,:,:,(0,dim-2)] *= 2.0self.weight_g = torch.ones(3,2,dim,dim-1,dim, dtype=torch.float)self.weight_g[:,:,:,(0,dim-2),:] *= 2.0self.weight_b = torch.ones(3,2,dim-1,dim,dim, dtype=torch.float)self.weight_b[:,:,(0,dim-2),:,:] *= 2.0self.relu = torch.nn.ReLU()def forward(self, LUT):dif_context = LUT.LUT_en[:,:-1,:,:,:] - LUT.LUT_en[:,1:,:,:,:]dif_r = LUT.LUT_en[:,:,:,:,:-1] - LUT.LUT_en[:,:,:,:,1:]dif_g = LUT.LUT_en[:,:,:,:-1,:] - LUT.LUT_en[:,:,:,1:,:]dif_b = LUT.LUT_en[:,:,:-1,:,:] - LUT.LUT_en[:,:,1:,:,:]tv = torch.mean(torch.mul((dif_r ** 2),self.weight_r)) + torch.mean(torch.mul((dif_g ** 2),self.weight_g)) + torch.mean(torch.mul((dif_b ** 2),self.weight_b)) mn = torch.mean(self.relu(dif_r)) + torch.mean(self.relu(dif_g)) + torch.mean(self.relu(dif_b)) \+ torch.mean(self.relu(dif_context))return tv, mn

這個沒有特別需要講解的，基本上是照著論文給出的公式將其翻譯成代碼，tv代表平滑性損失，mn代表單調(diào)性損失，因此這個類會同時輸出兩個損失，至于平滑損失中的w正則會在后續(xù)的訓練中看到。

train.py 文件

存放著跟訓練相關(guān)的代碼。以一個epoch的一個batch的一次iteration為例：

for epoch in range(opt.epoch, opt.n_epochs):mse_avg = 0psnr_avg = 0Generator_bias.train()Generator_context.train()for i, batch in enumerate(dataloader):# Model inputsreal_A = Variable(batch["A_input"].type(Tensor))real_B = Variable(batch["A_exptC"].type(Tensor))# ------------------ #  Train Generators# ------------------optimizer_G.zero_grad()fake_B, weights_norm = generator_train(real_A)# Pixel-wise lossmse = criterion_pixelwise(fake_B, real_B)tv_enhancement, mn_enhancement = TV4(LUT_enhancement)tv_cons = tv_enhancementmn_cons = mn_enhancement# loss = mseloss = mse + opt.lambda_smooth * (weights_norm + tv_cons) + opt.lambda_monotonicity * mn_conspsnr_avg += 10 * math.log10(1 / mse.item())mse_avg += mse.item()loss.backward()optimizer_G.step()

real_A 和real_B分別是增強前圖像和增強后的HQ，generator_train是根據(jù)LUT生成圖像的過程，實現(xiàn)如下所示：

def generator_train(img):context = Generator_context(img)pred = Generator_bias(img)context = context.new(context.size())context = Variable(context.fill_(0).type(Tensor))pred = pred.squeeze(2).squeeze(2)combine = torch.cat([context,img],1)gen_A0 = LUT_enhancement(combine)weights_norm = torch.mean(pred ** 2)combine_A = img.new(img.size())for b in range(img.size(0)):combine_A[b,0,:,:] = pred[b,0] * gen_A0[b,0,:,:] + pred[b,1] * gen_A0[b,1,:,:] + pred[b,2] * gen_A0[b,2,:,:] + pred[b,9]combine_A[b,1,:,:] = pred[b,3] * gen_A0[b,0,:,:] + pred[b,4] * gen_A0[b,1,:,:] + pred[b,5] * gen_A0[b,2,:,:] + pred[b,10]combine_A[b,2,:,:] = pred[b,6] * gen_A0[b,0,:,:] + pred[b,7] * gen_A0[b,1,:,:] + pred[b,8] * gen_A0[b,2,:,:] + pred[b,11]return combine_A, weights_norm

這里的Generator_context是前面剛講到的context encoder網(wǎng)絡結(jié)構(gòu)，Generator_bias是前面講到的param encoder的網(wǎng)絡結(jié)構(gòu)，
針對此處博主提出一個疑問，作者在代碼中初始化了一個空的context將其與輸入圖像進行cat得到了一個4D圖像，但如果輸入為空的如何進行訓練context encoder？
接下來會將combine對應的4D圖像輸入LUT_enhancement模塊進行增強處理。需要注意的是，該模塊僅使用一條4DLUT（通過4D插值生成輸出gen_A0），而非論文中所述的三條4DLUT。
最終計算param的L2正則norm損失并執(zhí)行LUT融合。此處將先前預測的12個參數(shù)用于輸出融合，融合過程會利用LUT的每個維度。因此是需要 12=3 * 1 * 3 + 1 * 3，3個通道乘以1條LUT乘以3個通道代表的權(quán)重項加上1條LUT乘以3個通道代表的偏置項。若這么理解的話，則與論文中給出的公式不一致，論文中給出的權(quán)重的個數(shù)是$3?N_{lut}^2$，我們這里看代碼只能理解為$3?3?N_{lut}$。

之后是計算損失的過程：

        # Pixel-wise lossmse = criterion_pixelwise(fake_B, real_B)tv_enhancement, mn_enhancement = TV4(LUT_enhancement)tv_cons = tv_enhancementmn_cons = mn_enhancement# loss = mseloss = mse + opt.lambda_smooth * (weights_norm + tv_cons) + opt.lambda_monotonicity * mn_cons

包含mse損失和正則、單調(diào)損失，正則損失使用的是我們前面講到的TV_4D類，最后進行權(quán)重進行加權(quán)融合即可。

3、總結(jié)

代碼實現(xiàn)核心的部分講解完畢，但作者開源的4DLUT代碼還是不完整和存在很多疑問點的，目前可以看到很多點與實際論文表述不一致，上面提到的疑問點總結(jié)有：

1. 4DLUT的初始化中關(guān)于context的維度個數(shù)，代碼中顯示是2，而非跟顏色通道一樣大，這樣如何進行部署。
2. 4DLUT訓練中關(guān)于context的初始化，如果在計算完context后不使用，而是使用一個0初始化的context，如何將context信息加入到4DLUT中。
3. 基礎4DLUT的個數(shù)，不是3條而是單獨的1條。
4. 4DLUT中param encoder預測param的個數(shù)，跟論文中給出的公式不一致，論文中給出的weight權(quán)重的個數(shù)是$3?N_{lut}^2$，但實際是$3?3?N_{lut}$。

4DLUT雖然利用context解決了圖像中局部區(qū)域增強的問題，但是有一個問題是利用4DLUT來推理，加入了更多的計算量和更多的限制，這可能會減少其適用的場景，實用性相比較3DLUT有所降低。

---

感謝閱讀，歡迎留言或私信，一起探討和交流。
如果對你有幫助的話，也希望可以給博主點一個關(guān)注，感謝。