本文是對CLUT技術的代碼講解，原文解讀請看CLUT文章講解。

1、原文概要

CLUT利用矩陣在保持3DLUT映射能力的前提下顯著降低了參數量。整體流程如下所示。

整體還是基于3D-LUT的框架，只不過添加了一個壓縮自適應的變換矩陣。作者使用的損失函數在3DLUT的基礎上額外添加了一個余弦相似度的損失。

2、代碼結構

代碼整體結構如下：

核心代碼是models.py與LUT.py文件。

3 、核心代碼模塊

model.py 文件

1. CLUTNet類

這里是網絡的整體實現，其定義了backbone、classifier、CLUT。

class CLUTNet(nn.Module): def __init__(self, nsw, dim=33, backbone='Backbone', *args, **kwargs):super().__init__()self.TrilinearInterpolation = TrilinearInterpolation()self.pre = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])self.backbone = eval(backbone)()last_channel = self.backbone.last_channelself.classifier = nn.Sequential(nn.Conv2d(last_channel, 128,1,1),nn.Hardswish(inplace=True),nn.Dropout(p=0.2, inplace=True),nn.Conv2d(128, int(nsw[:2]),1,1),)nsw = nsw.split("+")num, s, w = int(nsw[0]), int(nsw[1]), int(nsw[2])self.CLUTs = CLUT(num, dim, s, w)def fuse_basis_to_one(self, img, TVMN=None):mid_results = self.backbone(self.pre(img))weights = self.classifier(mid_results)[:,:,0,0] # n, numD3LUT, tvmn_loss = self.CLUTs(weights, TVMN)return D3LUT, tvmn_loss    def forward(self, img, img_org, TVMN=None):D3LUT, tvmn_loss = self.fuse_basis_to_one(img, TVMN)img_res = self.TrilinearInterpolation(D3LUT, img_org)return {"fakes": img_res + img_org,"3DLUT": D3LUT,"tvmn_loss": tvmn_loss,}

前向中給出了計算過程，首先圖像經過backbone計算中間結果，然后經過classifer得到CLUT的輸入，最后給到CLUT變換得到實際使用的3DLUT。

2. CLUT類

定義了CLUT的計算過程，講解中提到了有3個主要參數，num代表LUT的條數，s和w是壓縮的參數。

class CLUT(nn.Module):def __init__(self, num, dim=33, s="-1", w="-1", *args, **kwargs):super(CLUT, self).__init__()self.num = numself.dim = dimself.s,self.w = s,w = eval(str(s)), eval(str(w))# +: compressed;  -: uncompressedif s == -1 and w == -1: # standard 3DLUTself.mode = '--'self.LUTs = nn.Parameter(torch.zeros(num,3,dim,dim,dim))elif s != -1 and w == -1:  self.mode = '+-'self.s_Layers = nn.Parameter(torch.rand(dim, s)/5-0.1)self.LUTs = nn.Parameter(torch.zeros(s, num*3*dim*dim))elif s == -1 and w != -1: self.mode = '-+'self.w_Layers = nn.Parameter(torch.rand(w, dim*dim)/5-0.1)self.LUTs = nn.Parameter(torch.zeros(num*3*dim, w))else: # full-version CLUTself.mode = '++'self.s_Layers = nn.Parameter(torch.rand(dim, s)/5-0.1)self.w_Layers = nn.Parameter(torch.rand(w, dim*dim)/5-0.1)self.LUTs = nn.Parameter(torch.zeros(s*num*3,w))print("n=%d s=%d w=%d"%(num, s, w), self.mode)def reconstruct_luts(self):dim = self.dimnum = self.numif self.mode == "--":D3LUTs = self.LUTselse:if self.mode == "+-":# d,s  x  s,num*3dd  -> d,num*3dd -> d,num*3,dd -> num,3,d,dd -> num,-1CUBEs = self.s_Layers.mm(self.LUTs).reshape(dim,num*3,dim*dim).permute(1,0,2).reshape(num,3,self.dim,self.dim,self.dim)if self.mode == "-+":# num*3d,w x w,dd -> num*3d,dd -> num,3dddCUBEs = self.LUTs.mm(self.w_Layers).reshape(num,3,self.dim,self.dim,self.dim)if self.mode == "++":# s*num*3, w  x   w, dd -> s*num*3,dd -> s,num*3*dd -> d,num*3*dd -> num,-1CUBEs = self.s_Layers.mm(self.LUTs.mm(self.w_Layers).reshape(-1,num*3*dim*dim)).reshape(dim,num*3,dim**2).permute(1,0,2).reshape(num,3,self.dim,self.dim,self.dim)D3LUTs = cube_to_lut(CUBEs)return D3LUTsdef combine(self, weights, TVMN): # n,numdim = self.dimnum = self.numD3LUTs = self.reconstruct_luts()if TVMN is None:tvmn_loss = 0else:tvmn_loss = TVMN(D3LUTs)D3LUT = weights.mm(D3LUTs.reshape(num,-1)).reshape(-1,3,dim,dim,dim)return D3LUT, tvmn_lossdef forward(self, weights, TVMN=None):lut, tvmn_loss = self.combine(weights, TVMN)return lut, tvmn_loss

mode這里是調整壓縮的模式，當然我們需要的是完全壓縮的版本，即mode==“++”，可以看到首先會對w_layers與self.LUTs矩陣乘，后續在跟s_layers進行矩陣乘，這與講解相對應。

utils/LUT.py 文件

1. cube_to_lut函數

此函數在CLUT類的前向完成處理最后會調用到。

def cube_to_lut(cube): # (n,)3,d,d,dif len(cube.shape) == 5:to_shape = [[0,2,3,1],[0,2,1,3],]else:to_shape = [[1,2,0],[1,0,2],]if isinstance(cube, torch.Tensor):lut = torch.empty_like(cube)lut[...,0,:,:,:] = cube[...,0,:,:,:].permute(*to_shape[0])lut[...,1,:,:,:] = cube[...,1,:,:,:].permute(*to_shape[1])lut[...,2,:,:,:] = cube[...,2,:,:,:]else:lut = np.empty_like(cube)lut[...,0,:,:,:] = cube[...,0,:,:,:].transpose(*to_shape[0])lut[...,1,:,:,:] = cube[...,1,:,:,:].transpose(*to_shape[1])lut[...,2,:,:,:] = cube[...,2,:,:,:]return lut

通過CLUT類我們可以看到送入到該函數的輸入的shape是(num,3,self.dim,self.dim,self.dim)，因為shape的長度為5，to_shape是[0,2,3,1]以及[0,2,1,3]，也就是說實際的lut是調換通道順序的cube變量。

3、總結

代碼實現核心的部分講解完畢，跟以往最不同的部分就在于這個CLUT的計算矩陣，把這部分看明白就行。

---

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