深度学习图像压缩基础及经典讲解

目的

  在极低的码率下，保证图像重建的质量。
$$\mathcal{L}=R+\lambda D$$
  通常L是损失函数，R是码率，D是重建后与原图的失真，λ 是控制率失真权衡的系数。

Baseline

  基于AutoEnconder的一个过程，原始的图像x(Input Image)通过一个映射ga(通常是卷积构成的网络)变换到一个LatentSpace潜在空间，转换为了潜在表示 y。潜在空间便于量化和熵编码，量化和熵编码后，把比特流发送到解码器端解码。

Hyperprior

  由Balle提出的框架，由于Baseline在Entropy encoding中需要知道tensor的具体的值以及概率值才能编成码流文件。出现的概率越大分配码长越短，才能保证占用储存空间小。Baseline中也有概率模型并且在训练过程中每一个值都有概率表对应，但是一旦训练好了就变成固定的了，无论输入x是什么图像都是固定的，导致会出现了很大的冗余。
  所以Balle等提出了Hyperprior模型是一个双层的，假设latent representation是一个独立同分布的，它到底服从了一个什么分布，根据香农公式信息熵：
$$\mathrm{H}(\mathrm{X})=-\sum _{\mathrm{x}\in \chi }\mathrm{p}(\mathrm{x})\log \mathrm{p}(\mathrm{x})$$

  我们知道了分布就很好求熵，我们只要优化熵就行了。我们假设latent representation服从一个具体的概率分布，例如高斯分布、拉普拉斯分布等。分布在训练的时候可以是连续的，编码是离散的。在Hyperprior中假设服从了均值为0，标准差为δ的单高斯分布，Hyperprior就是多加了一层，这是一个自适应的过程，每输入一个x潜在空间概率分布是可以学习的，不同图像学习参数不一样。

Joint

  这里面把单高斯改为了均值为μ的，均值也是可以学习的，还加了Context Moudle，具有因果关系，通常用掩膜卷积。

  通过对卷积核进行 Mask操作，掩盖了未解码点的数值，保证了自回归模型预测的当前的像素点参数仅来自于前面已经解码的点，而不取决于未解码点。通过这种Mask 掩膜，遮蔽卷积核下面和右边的权重，这种卷积核与特征图进行卷积的时候，可以可知卷积得到的结果与“未来”的数据无关。自回归模型的问题在于存在严格是时序关系，只能先得到前面的点才能得到当前点的信息，而前面的点也只能依靠更前面的点得到，即表现为在解码的时候，原始的解码方式具有并行性质，而自回归则是串行顺序。

ProposedModel

  使用了混合高斯，在数学中混合高斯可以拟合任意的分布。

VAE(变分自动编码器)

  了解超先验首先要了解变分自动编码器理论基础。
参考：
  变分自编码器
  生成模型VAE
  VAE

极大似然估计

  什么是极大似然估计，一个猎人和一个从未打过猎的普通人一起去森林打猎，两人看到一只兔子，一声枪响，兔子死了。问是谁打死的兔子？关于这个例子，人们的猜测是猎人打死的兔子。因为猎人经常打猎，而普通人没有打过猎。这种猜测就是“最大似然原理”
  1.假设我们的抽样是理想正确的；2.概率大的事件在一次观测中更容易发生；3.在一次观测中发生了的事件其概率应该大；
  例如：推测全国人民身高，被抽取的n个人他们的概率一定大，概率大的事件更容易发生。所以我们可以认为，被抽取的n个人的身高，就是全国人们大概率的身高，或者说全国人们的身高，大概率都在被抽取的n个人身高中。这样我就可以用这n个人的身高来代替全国人的身高。
  现在我知道抽取的这n个人的概率一定是最大的了，且假设全国人身高服从正态分布，想知道这个正态分布就需要求它的均值和方差。这里均值跟方差就是我们想求解的未知参数，这个均值和方差我们就可以用极大似然估计来求。
  正态分布的概率密度函数为：
$$p(x;\mu ,{\sigma }^2)=\frac{1}{\sqrt{2\pi }\delta }{e}^{-\frac{(x-\mu {)}^2}{2{\sigma }^2}}$$
  求解步骤为：
step1 计算似然函数:
$$L(\mu ,{\sigma }^2)=\prod _{i=1}^n\frac{1}{\sqrt{2\pi }\delta }{e}^{-\frac{(x_i-\mu {)}^2}{2{\sigma }^2}}$$
step2 似然函数取对数：
$$lnL(\mu ,{\sigma }^2)=-\frac{n}{2}ln(2\pi )-\frac{n}{2}ln{\sigma }^2-\frac{1}{2{\sigma }^2}\sum _{i=1}^n(x_i-\mu {)}^2$$
step3 求最值对应的参数：
  我们令:$$\frac{\partial }{\partial \mu }lnL(\mu ,{\sigma }^2)=0$$$$\frac{\partial }{\partial {\sigma }^2}lnL(\mu ,{\sigma }^2)=0$$
step4得到估计参数：
$$\begin{array}{rl}& \tilde{\mu }=\frac{1}{n}\sum _{i=1}^n{x}_i\\ & \\ & {\tilde{\sigma }}^2=\frac{1}{n}\sum _{i=1}^n(x_i-\bar{x}{)}^2\end{array}$$
  在机器学习和深度学学习中由于给定的样本总是有限，所以我们希望用有限的样本数据的经验分布来代替真实数据的分布。概率生成模型的本质也就是通过有限的样本数据进行训练，得到生成模型，生成模型生成的数据最后全都符合真实数据的分布。

EM算法

  极大似然估计就是利用极大似然的原理去解决概率生成模型的参数估计问题，像之前文章中举得对全国人们身高进行估计的例子，就是假设全国人民身高服从高斯分布，已知条件是抽样的n个人的身高，然后利用极大估计去求解这个高斯分布的均值和方差，然后得出答案。
  结果和参数相互对应的时候，似然和概率在数值上是相等的，如果用 θ 表示环境对应的参数，x 表示结果，那么概率可以表示为$P(x|\theta )$ 。$P(x|\theta )$ 是条件概率的表示方法，θ 是前置条件，理解为在 θ 的前提下，事件 x 发生的概率，相对应的似然可以表示为$L(\theta |x)$
  求使得出现该组样本的概率最大的θ值:
$$\hat{\theta }=\arg \underset{\theta }{\max }l(\theta )=\arg \underset{\theta }{\max }\prod _{i=1}^{N}p(x_i\mid \theta )$$ $$H(\theta )=\ln l(\theta )$$ $$\hat{\theta }=\arg \underset{\theta }{\max }H\left(\theta \right)=\arg \max \ln l\left(\theta \right)=\arg \underset{\theta }{\max }\sum _{i=1}^N\ln p\left(x_i|\theta \right)$$ $$\frac{dl(\theta )}{d\theta }=0\text{或者等价于}\frac{dH(\theta )}{d\theta }=\frac{d\ln l(\theta )}{d\theta }=0$$
  其实EM算法跟极大似然估计一样，也是用来解决参数估计问题的。EM算法的目的是解决具有隐变量的混合模型的参数估计问题。注意两个点：隐变量和混合模型
  极大似然估计可以解决的全国人们身高估计问题是单高斯模型的问题。单高斯模型还是比较简单，实际上还有高斯混合模型（GMM）。
  下图这个例子的数据集被称为“老忠实间歇喷泉”数据集，由美国黄石国家公园的老忠实间歇喷泉的272次喷发的测量数据组成。每条测量记录包括喷发持续了几分钟（横轴）和距离下次喷发间隔了几分钟（纵轴）。我们看到数据集主要聚集在两大堆中，一个简单的高斯分布不能描述这种结构，而两个高斯分布的线性叠加可以更好地描述这个数据集的特征。

  左图，我们用一个高斯模型对样本来做分析，显然，一个高斯分布是不合适的。一般来说越靠近椭圆中心样本出现的概率越大，这是由概率密度函数决定的，但是这个高斯分布的椭圆中心的样本量却极少。显然样本服从单高斯分布的假设并不合理。单高斯模型无法产生这样的样本。基于此，我们引入了高斯混合模型，我们发现用两个高斯模型通过一定的权重形成的混合高斯模型可以产生这样的样本，如上面右图。

  高斯混合模型的本质就是融合几个单高斯模型，来使的模型更加复杂，从而产生更复杂的样本。理论上，如果某个混个高斯模型融合的高斯模型个数足够多，他们之间的权重设定的足够合理，这个混合模型可以拟合任意分布的样本。
  设有随机变量X，则混合高斯模型可以用下式表示：
$$\mathrm{p}(\mathbf{x})=\sum _{\mathrm{k}=1}^{\mathrm{K}}{\pi }_{\mathrm{k}}\mathcal{N}(\mathbf{x}|{\mathbf{\mu }}_{\mathrm{k}},{\mathbf{\Sigma }}_{\mathrm{k}})$$
  其中$\mathcal{N}(\mathbf{x}|{\mathbf{\mu }}_{\mathrm{k}},{\mathbf{\Sigma }}_{\mathrm{k}})$称为混合模型中的第k个分量。如前面图中的例子，有两个聚类，可以用两个二维高斯分布来表示，那么分量数K=2。${\pi }_k$是混合系数，且满足：$$\sum _{\mathrm{k}=1}^{\mathrm{K}}{\pi }_{\mathrm{k}}=1$$$$0\le {\pi }_{\mathrm{k}}\le 1$$
  实际上，可以认为${\pi }_k$ 就是每个分量$\mathcal{N}(\mathbf{x}|{\mathbf{\mu }}_{\mathrm{k}},{\mathbf{\Sigma }}_{\mathrm{k}})$的权重。
  例如图上的例子，很明显有两个聚类，可以定义K = 2. 那么对应的GMM形式如下：
$$\mathrm{p}(\mathbf{x})={\pi }_1\mathcal{N}(\mathbf{x}|{\mathbf{\mu }}_1,{\mathbf{\Sigma }}_1)+{\pi }_2\mathcal{N}(\mathbf{x}|{\mathbf{\mu }}_2,{\mathbf{\Sigma }}_2)$$
  上式中未知的参数有六个：$({\pi }_1,{\mathbf{\mu }}_1,{\mathbf{\Sigma }}_1;{\mathbf{\pi }}_2,{\mathbf{\mu }}_2,{\mathbf{\Sigma }}_2).$ 之前提到GMM聚类时分为两步，第一步是随机地在这K个分量中选一个，每个分量被选中的概率即为混合系数${\pi }_k$。可以设定${\pi }_1={\pi }_2=0.5\text{,}$表示每个分量被选中的概率是0.5，即从中抽出一个点，这个点属于第一类的概率和第二类的概率各占一半。但实际应用中事先指定${\pi }_k$ 的值是很笨的做法，当问题一般化后，会出现一个问题：当从图2中的集合随机选取一个点，怎么知道这个点是来自$N(\mathbf{x}|{\mathbf{\mu }}_1,{\mathbf{\Sigma }}_1)\text{还是N}(\mathbf{x}|{\mathbf{\mu }}_2,{\mathbf{\Sigma }}_2)$呢？换言之怎么根据数据自动确定${\pi }_1和{\pi }_2$的值呢？这就是GMM参数估计的问题。要解决这个问题，可以使用EM算法。通过EM算法，我们可以迭代计算出GMM中的参数：$({\pi }_{\mathrm{k}},{\mathbf{x}}_{\mathrm{k}},{\mathbf{\Sigma }}_{\mathrm{k}}).$这是高斯混合模型的参数，如果我们按照前面单高斯模型求解参数，则我们要用似然函数来求解，此混合模型的对数似然函数为：$$\log L(u,\Sigma ,\pi )=\sum _{i=1}^N\log \sum _{k=1}^k{\pi }_{k}N(x_i|{\mu }_k,{\Sigma }_k)$$
用MLE求参数θ为：$${\hat{\theta }}_{mLE}=\arg \max \left\{\sum _{i=1}^N\log \left[\sum _{k=1}^k{\pi }_{k}N(X_i|{\mu }_k,{\Sigma }_k)\right]\right\}$$  根据上图，我们可以看到，如果我们用MLE来求解高斯混合模型的话，我们很难求得，因为上式中log后面有一个求和并且高斯分布也是多维的，这样我们是无法通过求导得到参数。
  期望最大算法的目的是解决具有隐变量的混合模型的参数估计 (极大似然估计)。MLE 对$p(x|\theta )$ 参数的估计记为： ${\theta }_{MLE}=argmax\log p(x|\theta )$ 。EM 算法对这个问题的解决方法是。
采用迭代的方法：
$${\theta }^{t+1}=argmax{\int }_z\log [p(x,z|\theta )]p(z|x,{\theta }^t)dz={\mathbb{E}}_{z|x,{\theta }^t}\left[\log p(x,z|\theta )\right]$$
  这个公式包含了迭代的两步：
E. step: 计算 $\log p(x,z|\theta )$ 在概率分布 $p(z|x,{\theta }^t)$ 下的期望
M.step: 计算使这个期望最大化的参数得到下一个 EM 步骤的输入
  EM算法具有收敛性,EM算法是逐步迭代来求解使期望最大的参数θ的$$\log p(x|{\theta }^t)\le \log p(x|{\theta }^{t+1})$$

VAE

  AE(自动编码器)

  AE主要有encoder和decoder两个部分组成，其中encoder和decoder都是神经网络。其中encoder负责将高维输入转换为低维的code，decoder负责将低维的code转换为高维的输入，其中输出要跟输入尽可能的相似，最好是能完全一样。
  图的中间部分，就代表了code，我们经过训练，最终得到了一个6维的code，且每个维度有一个具体的值代表。因为AE就是想生成跟输入相同的图像，所以AE是不会产生新的图像的，就像上面第二个图，最后生成的这个大叔，不会是一个闭嘴不笑，红头发的大叔。
  再举一个例子：

  如上图，经过训练我们的AE已经能还原这两张图片。接下来，我们在code空间上，两张图片的编码点中间处取一点，然后将这一点交给解码器，我们希望新的生成图片是一张清晰的图片（类似3/4全月的样子）。但是，实际的结果是，生成图片是模糊且无法辨认的乱码图。
  为什么我们在code的中间取一点，无法生成3/4月亮的图呢？
还是因为，AE就是要生成跟输入尽可能相同的图，类似与压缩与解压缩的过程，你想让AE生成一个相似但全新的图，它是做不到的。
  那如果我就是想生成3/4月亮的图怎么办？有一种方法，我们可以引入噪声，使得图片的编码区域得到扩大，从而掩盖掉失真的空白编码点。如下图所示：

  如上图所示，现在在给两张图片编码的时候加上一点噪音，使得每张图片的编码点出现在绿色箭头所示范围内，于是在训练模型的时候，绿色箭头范围内的点都有可能被采样到，这样解码器在训练时会把绿色范围内的点都尽可能还原成和原图相似的图片。然后我们可以关注之前那个失真点，现在它处于全月图和半月图编码的交界上，于是解码器希望它既要尽量相似于全月图，又要尽量相似于半月图，于是它的还原结果就是两种图的折中（3/4全月图）。
  由此我们发现，给编码器增添一些噪音，可以有效覆盖失真区域。不过这还并不充分，因为在上图的距离训练区域很远的黄色点处，它依然不会被覆盖到，仍是个失真点。为了解决这个问题，我们可以试图把噪音无限拉长，使得对于每一个样本，它的编码会覆盖整个编码空间，不过我们得保证，在原编码附近编码的概率最高，离原编码点越远，编码概率越低。在这种情况下，图像的编码就由原先离散的编码点变成了一条连续的编码分布曲线，如下图所示。

  那么上述的这种将图像编码由离散变为连续的方法，就是变分自编码的核心思想。VAE中间的code不是具体的值，而是一种分布。如下图所示：

  VAE(变分自动编码器)
  VAE是从AE的基础上来的，其中我们知道，AE类似于一种压缩解压缩的过程，只能生成与输入图像尽可能一样的图像，而不会生成与输入图像相似但不同的图像，这是AE本身所限制的。在这个基础上，VAE中间的code是一个分布，从分布中采样然后输入到decoder中，VAE可以生成与输入图像相似但不同的图像。

  如下图所示，为VAE原文中给出的VAE的图模型，从VAE的图模型我们可以看到，z是服从标准正太分布的。

  如下图所示，为VAE原文中给出的VAE的结构图。其中分为左右两个部分，右边才是最终的VAE结构。右边比左边就多了一个从标准正太分布中采样的过程，称为重参数化技巧，后面会说到为什么要有这个重参数化技巧。

  首先，VAE是一个生成模型，准确的说它是深度概率生成模型。概率生成模型关心样本的分布，对样本分布本身进行建模。
  其次，VAE是隐变量模型，隐变量模型是存在隐变量z，其中GMM也是一个隐变量模型，GMM中也存在隐变量z。VAE的生成样本是由隐变量z生成的，如VAE原文中的图模型所示。
  介绍一下各种符号表示的意义：
$$X：\text{观测数据}{q}_{\varphi }(z\mid x):\text{encoder(近似后验)}$$$$z：\text{隐变量}{p}_{\theta }(z|x)：\text{真实后验}$$$$\varphi ：\text{变分参数}{p}_{\theta }(x,z)：\text{生成模型}$$$$\theta ：\text{生成参数}{p}_{\theta }(x|z)：\text{decoder}$$
  我们知道VAE是深度概率生成模型。概率生成模型关心样本的分布，对样本分布本身进行建模。因此，样本分布为$P(X)$其中$X=x1,x2,x3,...x{n}_{\text{。}}$，目前我们已知的只有一些来自分布P(X)的样本，且我们并不知道P(X)是什么分布，现在我们只能根据一些已知的样本去求样本分布P(X)。显而易见，我们是无法求解的。这样我们不妨转换一下思路，既然根据已知样本无法直接求解P(X)，那么我们做出合理的假设，假设样本X,是由隐变量z产生的。这样一来，$P(x)=\underset{z}{\int }P(z)P\left(x|z\right)dz$,这样一来P(X)就有了一个基本的结构。
  关于隐变量z的分布$P(Z)$,我们也做出了假设，假设它是服从标准正太分布的，即$P(Z)\sim N(0,I)$。这里$P(Z)$不一定非要这样假设，你也可以假设它是其他的分布，那为什么这里要如此假设呢？ 我觉得有一下两点理由：
   1.标准正太分布足够简单，并且理论上来说，标准正太分布经过变换可以拟合任意复杂的分布。
   2.中心极限定理，也称为大数定理。也就是说只要样本足够大，我们都可以看作其是一个正太分布。
   关于$P(X|Z)$我们也假设它服从高斯分布，即$P(X|Z)\sim N(\mu ,\Sigma )$,理由如上。
   有了合理的假设，我们继续推导$P(X)$,这里用到的思想是极大似然的思想，简单来说就是在理想情况下，现有的样本的概率一定大，概率大的样本更容易被获取。凭借极大似然的思想，我们现有的样本的概率一定大，我们就最大化这个概率$P(X)$,其中X为所有现有的样本的集合。
  通常我们会将公式改写成log函数的形式，因为公式推导过程中会出现累乘，用log函数可以将累乘转换为累加，方面推导。并且整体转换为log形式，并不改变其相对大小。在极大似然的计算中我们也是经常使用对数似然来计算。
  则我们改写成对数似然的形式：
$$logP(X)=logP(X,Z)-logP(Z|X)=log\frac{P(X,Z)}{q(z)}-log\frac{P(Z|X)}{q(z)}$$
此时，分别对公式的左右两边求期望$E_{q(z)}$
其中，左边： $$\int q(z)logP(X)dz=logP(X)\int q(z)dz=logP(X)$$
右边：
$$\begin{array}{l}\int q(z)log\frac{P(X,Z)}{q(z)}dz-\int q(z)log\frac{P(Z|X)}{q(z)}dz=ELBO+KL(q(z)||P(Z|X)))\end{array}$$
  可以看到，$$logP(X)=ELBO-KL(q(z)||P(Z|X))$$

  其中，等式右边最后一项为KL>=0,当$q(z)=P(Z|X)\text{时，KL=}0$。则$LogP(X)>=ELBO$,我们称ELBO为变分下界。此时，我想最大化P(X), 其实可以转化成最大化ELBO。
  其中， ELBO也可以继续推导：
$$\begin{array}{rl}& ELBO=\int q(z)log\frac{P(X,Z)}{q(z)}dz=\int q(z)log\frac{P(X|Z)P(Z)}{q(z)}dz=\\ & \int q(z)logP(X|Z)dz+\int q(z)log\frac{P(Z)}{q(z)}dz=E_{q(z)}logP(X|Z)-\\ & KL(q(z)||P(Z))\end{array}$$
  可以看到，ELBO被分解成两项，第一项属于重构，第二项是一个正则化项。我们将ELBO作为损失函数来训练VAE，因为是最大化ELBO。
  至此，公式算是基本推导完成。那么你可能疑问变分在哪里呢？?其实，变分就是我们引入的q(z),我们让q(z) 去拟合不可计算的P(Z|X)。那么为什么P(Z|X)不可计算呢?因为$$P(Z|X)=\frac{P(X,Z)}{P(X)}$$
  其中P(X)是我们要求的。这里进一步说，q(z)中的z是由样本得到的，所以q(z)=q(z|x)。我们可以将上面的q(z)都替换成q(z|x)带入公式，这里我们就不写了。
  我们用q(zlx)去拟合不可计算的P(ZIX)其实就是变分推断。假如是纯数学的方法，我们进行变分推断的话会有基于平均场的变分推断，SGVI等方法。这里，我们直接用神经网络来拟合，也就是encoder部分，输入样本x，得到均值和方差，然后从这个p(z)中采样一个z值，将z输入到decoder中，生成样本。

GaussianConditional(GSM)

  下面是来自compressAI的GaussianConditional类代码,来自文章。

class GaussianConditional(EntropyModel):
    r"""Gaussian conditional layer, introduced by J. Ballé, D. Minnen, S. Singh,
    S. J. Hwang, N. Johnston, in `"Variational image compression with a scale
    hyperprior" <https://arxiv.org/abs/1802.01436>`_.

    This is a re-implementation of the Gaussian conditional layer in
    *tensorflow/compression*. See the `tensorflow documentation
    <https://github.com/tensorflow/compression/blob/v1.3/docs/api_docs/python/tfc/GaussianConditional.md>`__
    for more information.
    """

    def __init__(
            self,
            scale_table: Optional[Union[List, Tuple]],
            *args: Any,
            scale_bound: float = 0.11,
            tail_mass: float = 1e-9,
            **kwargs: Any,
    ):
        super().__init__(*args, **kwargs)

        if not isinstance(scale_table, (type(None), list, tuple)):
            raise ValueError(f'Invalid type for scale_table "{type(scale_table)}"')

        if isinstance(scale_table, (list, tuple)) and len(scale_table) < 1:
            raise ValueError(f'Invalid scale_table length "{len(scale_table)}"')

        if scale_table and (
                scale_table != sorted(scale_table) or any(s <= 0 for s in scale_table)
        ):
            raise ValueError(f'Invalid scale_table "({scale_table})"')

        self.tail_mass = float(tail_mass)
        if scale_bound is None and scale_table:
            scale_bound = self.scale_table[0]
        if scale_bound <= 0:
            raise ValueError("Invalid parameters")
        self.lower_bound_scale = LowerBound(scale_bound)

        self.register_buffer(
            "scale_table",
            self._prepare_scale_table(scale_table) if scale_table else torch.Tensor(),
        )

        self.register_buffer(
            "scale_bound",
            torch.Tensor([float(scale_bound)]) if scale_bound is not None else None,
        )

    @staticmethod
    def _prepare_scale_table(scale_table):
        return torch.Tensor(tuple(float(s) for s in scale_table))

    def _standardized_cumulative(self, inputs: Tensor) -> Tensor:
        half = float(0.5)
        const = float(-(2 ** -0.5))
        # Using the complementary error function maximizes numerical precision.
        return half * torch.erfc(const * inputs)

    @staticmethod
    def _standardized_quantile(quantile):
        return scipy.stats.norm.ppf(quantile)

    def update_scale_table(self, scale_table, force=False):
        # Check if we need to update the gaussian conditional parameters, the
        # offsets are only computed and stored when the conditonal model is
        # updated.
        if self._offset.numel() > 0 and not force:
            return False
        device = self.scale_table.device
        self.scale_table = self._prepare_scale_table(scale_table).to(device)
        self.update()
        return True

    def update(self):
        multiplier = -self._standardized_quantile(self.tail_mass / 2)
        pmf_center = torch.ceil(self.scale_table * multiplier).int()
        pmf_length = 2 * pmf_center + 1
        max_length = torch.max(pmf_length).item()

        device = pmf_center.device
        samples = torch.abs(
            torch.arange(max_length, device=device).int() - pmf_center[:, None]
        )
        samples_scale = self.scale_table.unsqueeze(1)
        samples = samples.float()
        samples_scale = samples_scale.float()
        upper = self._standardized_cumulative((0.5 - samples) / samples_scale)
        lower = self._standardized_cumulative((-0.5 - samples) / samples_scale)
        pmf = upper - lower

        tail_mass = 2 * lower[:, :1]

        quantized_cdf = torch.Tensor(len(pmf_length), max_length + 2)
        quantized_cdf = self._pmf_to_cdf(pmf, tail_mass, pmf_length, max_length)
        self._quantized_cdf = quantized_cdf
        self._offset = -pmf_center
        self._cdf_length = pmf_length + 2

    def _likelihood(
            self, inputs: Tensor, scales: Tensor, means: Optional[Tensor] = None
    ) -> Tensor:
        half = float(0.5)

        if means is not None:
            values = inputs - means
        else:
            values = inputs

        scales = self.lower_bound_scale(scales)

        values = torch.abs(values)
        upper = self._standardized_cumulative((half - values) / scales)
        lower = self._standardized_cumulative((-half - values) / scales)
        likelihood = upper - lower

        return likelihood

    def forward(
            self,
            inputs: Tensor,
            scales: Tensor,
            means: Optional[Tensor] = None,
            training: Optional[bool] = None,
    ) -> Tuple[Tensor, Tensor]:
        if training is None:
            training = self.training
        outputs = self.quantize(inputs, "noise" if training else "dequantize", means)
        likelihood = self._likelihood(outputs, scales, means)
        if self.use_likelihood_bound:
            likelihood = self.likelihood_lower_bound(likelihood)
        return outputs, likelihood

首先看forward前向传播中代码：

outputs = self.quantize(inputs, "noise" if training else "dequantize", means)

  在这里inputs是我们输入给GaussianConditional的y=ga(x),means=none，这里主要是对y进行均匀量化(在训练期间叠加均匀噪声)，得到了y_hat。
  把y_hat和scales=h_s(z_hat)作为输入。

likelihood = self._likelihood(outputs, scales, means)

def _likelihood(
            self, inputs: Tensor, scales: Tensor, means: Optional[Tensor] = None
    ) -> Tensor:
        half = float(0.5)

        if means is not None:
            values = inputs - means
        else:
            values = inputs

        scales = self.lower_bound_scale(scales)

        values = torch.abs(values)
        upper = self._standardized_cumulative((half - values) / scales)
        lower = self._standardized_cumulative((-half - values) / scales)
        likelihood = upper - lower

        return likelihood

  具体看这里：

upper = self._standardized_cumulative((half - values) / scales)
lower = self._standardized_cumulative((-half - values) / scales)

    def _standardized_cumulative(self, inputs: Tensor) -> Tensor:
        half = float(0.5)
        const = float(-(2 ** -0.5))
        # Using the complementary error function maximizes numerical precision.
        return half * torch.erfc(const * inputs)

$const=-\frac{1}{\sqrt{2}}$,输入inputs是$\frac{0.5-x}{\delta }$,所以是torch.erfc($\frac{x-0.5}{\sqrt{2}\delta }$)。具体看一下torch.erfc：$$\mathrm{erfc}\left(\mathrm{x}\right)=1-\mathrm{erf}\left(\mathrm{x}\right)=\frac{2}{\sqrt{\pi }}{\int }_{\mathrm{x}}^{\infty }{\mathrm{e}}^{-{\eta }^2}\mathrm{d}\eta$$
它的数学意义是对应红色部分面积

所以分别是upper=0.5×torch.erfc($\frac{x-0.5}{\sqrt{2}\delta }$),lower=0.5×torch.erfc($\frac{x+0.5}{\sqrt{2}\delta }$)

likelihood = upper - lower

代表了红色区域面积，是在这一个区域里面的概率值。

EntropyBottleneck

              ┌───┐    y     ┌───┐  z  ┌───┐ z_hat      z_hat ┌───┐
        x ──►─┤g_a├──►─┬──►──┤h_a├──►──┤ Q ├───►───·⋯⋯·───►──┤h_s├─┐
              └───┘    │     └───┘     └───┘        EB        └───┘ │
                       ▼                                            │
                     ┌─┴─┐                                          │
                     │ Q │                                          ▼
                     └─┬─┘                                          │
                       │                                            │
                 y_hat ▼                                            │
                       │                                            │
                       ·                                            │
                    GC : ◄─────────────────────◄────────────────────┘
                       ·                 scales_hat
                       │
                 y_hat ▼
                       │
              ┌───┐    │
    x_hat ──◄─┤g_s├────┘
              └───┘

代码前向传播过程：

    def forward(
            self, x: Tensor, training: Optional[bool] = None
    ) -> Tuple[Tensor, Tensor]:
        if training is None:
            training = self.training

        if not torch.jit.is_scripting():
            # x from B x C x ... to C x B x ...
            perm = np.arange(len(x.shape))
            perm[0], perm[1] = perm[1], perm[0]
            # Compute inverse permutation
            inv_perm = np.arange(len(x.shape))[np.argsort(perm)]
        else:
            raise NotImplementedError()
            # TorchScript in 2D for static inference
            # Convert to (channels, ... , batch) format
            # perm = (1, 2, 3, 0)
            # inv_perm = (3, 0, 1, 2)

        x = x.permute(*perm).contiguous()
        shape = x.size()
        values = x.reshape(x.size(0), 1, -1)

        # Add noise or quantize

        outputs = self.quantize(
            values, "noise" if training else "dequantize", self._get_medians()
        )

        if not torch.jit.is_scripting():
            likelihood, _, _ = self._likelihood(outputs)
            if self.use_likelihood_bound:
                likelihood = self.likelihood_lower_bound(likelihood)
        else:
            raise NotImplementedError()
            # TorchScript not yet supported
            # likelihood = torch.zeros_like(outputs)

        # Convert back to input tensor shape
        outputs = outputs.reshape(shape)
        outputs = outputs.permute(*inv_perm).contiguous()

        likelihood = likelihood.reshape(shape)
        likelihood = likelihood.permute(*inv_perm).contiguous()

        return outputs, likelihood

  传入进来的值是z = self.h_a(torch.abs(y))，经过超先验编码后的z。
量化过程：

 outputs = self.quantize(
            values, "noise" if training else "dequantize", self._get_medians()
        )

  这一部分和上面类似输入Z_hat

likelihood, _, _ = self._likelihood(outputs)

    def _likelihood(
            self, inputs: Tensor, stop_gradient: bool = False
    ) -> Tuple[Tensor, Tensor, Tensor]:
        half = float(0.5)
        lower = self._logits_cumulative(inputs - half, stop_gradient=stop_gradient)
        upper = self._logits_cumulative(inputs + half, stop_gradient=stop_gradient)
        likelihood = torch.sigmoid(upper) - torch.sigmoid(lower)
        return likelihood, lower, upper

    def _logits_cumulative(self, inputs: Tensor, stop_gradient: bool) -> Tensor:
        # TorchScript not yet working (nn.Mmodule indexing not supported)
        logits = inputs
        for i in range(len(self.filters) + 1):
            matrix = getattr(self, f"_matrix{i:d}")
            if stop_gradient:
                matrix = matrix.detach()
            logits = torch.matmul(F.softplus(matrix), logits)

            bias = getattr(self, f"_bias{i:d}")
            if stop_gradient:
                bias = bias.detach()
            logits += bias

            if i < len(self.filters):
                factor = getattr(self, f"_factor{i:d}")
                if stop_gradient:
                    factor = factor.detach()
                logits += torch.tanh(factor) * torch.tanh(logits)
        return logits

matrix是一个权重矩阵，它来自于如下代码的初始化

filters: Tuple[int, ...] = (3, 3, 3, 3),

self.filters = tuple(int(f) for f in filters)

init_scale: float = 10,

self.init_scale = float(init_scale)

filters = (1,) + self.filters + (1,)
        scale = self.init_scale ** (1 / (len(self.filters) + 1))

 for i in range(len(self.filters) + 1):
            init = np.log(np.expm1(1 / scale / filters[i + 1]))
            matrix = torch.Tensor(channels, filters[i + 1], filters[i])
            matrix.data.fill_(init)

  bias是偏置矩阵如下

 bias = torch.Tensor(channels, filters[i + 1], 1)
            nn.init.uniform_(bias, -0.5, 0.5)
            self.register_parameter(f"_bias{i:d}", nn.Parameter(bias))

  对于X累积分布函数,即所有小于等于 x 的值出现概率的和。
$${\mathrm{F}}_{\mathrm{X}}\left(\mathrm{x}\right)=\mathrm{P}\left(\mathrm{X}\le \mathrm{x}\right)$$
  求到累积分布的上下界后，它们的差就是在这中间的概率了。

损失函数

KL散度

  相对熵（Relative Entropy），也叫 KL 散度 (Kullback-Leibler Divergence)，具有非负的特性。用于衡量两个分布之间距离的指标，用P分布近似Q的分布，相对熵可以计算这个中间的损失。
$$D_{\mathrm{KL}}(P\Vert Q)={\mathbb{E}}_{\mathrm{x}\sim P}\left[\log \frac{P(x)}{Q(x)}\right]={\mathbb{E}}_{\mathrm{x}\sim P}[\log P(x)-\log Q(x)]$$
  P往往表示样本的真实分布，Q表示模型所预测的分布

交叉熵

  交叉熵是用来衡量估计模型于真实概率分布之间差异情况的。如果一个随机变量 $X\sim p(x)$,q(x)为用于近似p(x)的概率分布，那么 和模型 q(x)之间的交叉熵 (cross entropy) 可以定义为:$$H(X,q)=H(x)+D(p||q)=-\sum _x\mathrm{p}(\mathrm{x})\mathrm{logq}(\mathrm{x})$$

  等于信息熵加上KL散度，可以看出估计模型 q(x) 和 真实模型 p(x)之间的差异 。

交叉熵与softmax

  分类问题中常用交叉熵作为模型的损失函数。样本标签 y 的值为1或者0可以看做是概率，而模型的输出是一个实数值，如何将这个实数值转换成概率呢？这就要用到 softmax 函数了（所以面试官会经常问为什么交叉熵要和 softmax 一起用）。假设模型输出为${\mathrm{y}}_1,{\mathrm{y}}_2,...,{\mathrm{y}}_{\mathrm{n}}$，经过 softmax 后的输出为：$$\mathrm{softmax}({\mathrm{y}}_{\mathrm{i}})=\frac{{\mathrm{e}}^{{\mathrm{y}}_{\mathrm{i}}}}{{\sum }_{\mathrm{j}=1}^{\mathrm{n}}{\mathrm{e}}^{{\mathrm{y}}_{\mathrm{i}}}}$$
  这样就把模型的输出也变成了一个概率分布，从而可以用交叉熵来计算预测值和真实值之间的距离了。

KL散度率失真

$${\mathbb{E}}_{x\sim p_x}{D}_{\mathrm{KL}}[q\parallel p_{\tilde{y}\mid x}]={\mathbb{E}}_{\mathbf{x}\sim p_x}{\mathbb{E}}_{\tilde{\mathbf{y}}\sim \mathbf{q}}[\log q(\tilde{\mathbf{y}}\mid \mathbf{x})\underset{\text{weighted distortion}}{\underbrace{-\log p_{\mathbf{x}\mid \tilde{\mathbf{y}}}(\mathbf{x}\mid \tilde{\mathbf{y}})}}\underset{\mathrm{rate}}{\underbrace{-\log p_{\tilde{\mathbf{y}}}(\tilde{\mathbf{y}})}}]+const.$$
  KL 散度的最小化相当于优化压缩模型的率失真性能，第一项的计算结果为零，第二项和第三项分别对应于加权失真和比特率。
  我们先了解一下熵的计算的期望形式：
$$H(P)={\mathbb{E}}_{x\sim P}[-\log P(x)]$$
  其中，x~P 表示用事件 x 服从概率分布 P，并且使用 P 来计算期望，同时熵一般用字母 H 表示。熵这一概念，是用来表征在遵循特定概率分布的事件事件 x 发生时，传递信息需要用到的理论平均最小编码大小。因此，只要我们知道发生事件的概率分布，我们就可以计算它的熵，如果我们不知道它的概率分布，那就没有办法计算熵。所以，对于不知道概率分布的事件，需要想办法来估计概率分布，从而计算熵。
  估计熵：
  假设在观察了东京一段时间的天气后，我大概得到了一个天气的概率分布Q，那么使用这个估计的概率分布Q，便可以计算出一个熵值（此时称为估计熵Estimated Entropy）为：
$$EstimatedEntropy={\mathbb{E}}_{x\sim Q}[-\log Q(x)]$$
如果 Q 非常接近真实发生的概率分布，那么上式计算的结果（估计熵）便能更精确的反应传递信息需要的最小编码大小。

  但是，这种计算估计熵的公式，存在两种不确定性。
  首先，x~Q 表示我们利用概率分布 Q 来计算期望，而 Q 是估计出来的，可能会与实际概率分布 P 相差很大。
  其次，计算最小编码时，是基于估计的 Q 来计算的概率分布的负对数 -log(Q），而这也不会是100%准确的。
  由此可见，Q 即影响期望，又影响最小编码估计，因此，计算出的估计熵并不能反应出什么，继而将估计熵和真实熵进行对比，也无法得出什么有意义的结论。
  交叉熵：
  因此，我们既然有了熵的理论了，可以基于熵的理论，将我们的计算的编码大小与理论最小的编码进行比较即可，而不去考虑期望，将Q的影响由双变量变为单变量。
  举个例子，假设在观察一段已发生的东京天气后，得到了天气发生的真实分布P ，我们可以使用概率分布 P 来计算真实的平均编码大小，而利用天气预报中的概率分布 Q，来计算未发生的事件所需要的编码大小。
  这便是基于概率事件 P 和 Q 之间的交叉熵：
$$H(P,Q)={\mathbb{E}}_{x\sim P}[-\log Q(x)]$$
  我们回过来看这个公式的第二项$-\log p_{\mathbf{x}|\tilde{\mathbf{y}}}(\mathbf{x}\mid \tilde{\mathbf{y}})$是一个条件自信息，由以$\tilde{\mathbf{y}}$为条件下的X的自信息，这是一个后验概率问题。通过解码得到的果，来推测原始图像x的像素。当两者像素值越接近的时候，其得到的后验概率越大，第二项的（即自信息）越小。因此等价于图像编码中的失真项。
  第三项，$-\log {\mathrm{p}}_{\tilde{\mathrm{y}}}(\tilde{\mathrm{y}})$在信息论中即表示$\tilde{\mathbf{y}}$其对应了编码框架中的码率估计模型或者说由熵编码后待编码点的码率大小。

  率失真代码：

class RateDistortionLoss(nn.Module):
    """Custom rate distortion loss with a Lagrangian parameter."""

    def __init__(self, lmbda=0.01, metric="mse", return_type="all"):
        super().__init__()
        if metric == "mse":
            self.metric = nn.MSELoss()
        elif metric == "ms-ssim":
            self.metric = ms_ssim
        else:
            raise NotImplementedError(f"{metric} is not implemented!")
        self.lmbda = lmbda
        self.return_type = return_type

    def forward(self, output, target):
        N, _, H, W = target.size()
        out = {}
        num_pixels = N * H * W

        out["bpp_loss"] = sum(
            (torch.log(likelihoods).sum() / (-math.log(2) * num_pixels))
            for likelihoods in output["likelihoods"].values()
        )
        if self.metric == ms_ssim:
            out["ms_ssim_loss"] = self.metric(output["x_hat"], target, data_range=1)
            distortion = 1 - out["ms_ssim_loss"]
        else:
            out["mse_loss"] = self.metric(output["x_hat"], target)
            distortion = 255**2 * out["mse_loss"]

        out["loss"] = self.lmbda * distortion + out["bpp_loss"]
        if self.return_type == "all":
            return out
        else:
            return out[self.return_type]

  误差选择均方误差，失真将以均方误差的形式度量。看前向传播部分，我们可以看出，传入的参数output, target。output里面存了重建的图片，y的概率分布，z的概率分布，target是原始图像。
  主要看这一部分：

out["bpp_loss"] = sum(
            (torch.log(likelihoods).sum() / (-math.log(2) * num_pixels))
            for likelihoods in output["likelihoods"].values()
        )

  python里面log()默认是e为底数的，所以这里面使用了对数的换底公式，换成了以2为底。这段代码计算了平均比特率，公式如下,num是像素总数，所有概率取对数再相加取平均就是平均比特率。
$$\sum [-\sum {\log }_{2}p]\cdot \frac{1}{num}$$
  比特率部分看完了我们来看失真部分这一段代码。

out["mse_loss"] = self.metric(output["x_hat"], target)
            distortion = 255**2 * out["mse_loss"]

  metric是均方误差，是重建后的图像和原始图像的均方误差，MSE计算了每个像素的预测值与目标值之间的平方差的平均值，在图像处理中，通常将像素值表示为介于0到255之间的整数，乘积255的平方代表了表示整个图像的像素值平均差异，标度可以使失真更易于理解。
  然后损失函数：

 out["loss"] = self.lmbda * distortion + out["bpp_loss"]

$$\mathcal{L}=R+\lambda D$$

GaussianConditional(GMM？)

  来自文章

              ┌───┐    y     ┌───┐  z  ┌───┐ z_hat      z_hat ┌───┐
        x ──►─┤g_a├──►─┬──►──┤h_a├──►──┤ Q ├───►───·⋯⋯·───►──┤h_s├─┐
              └───┘    │     └───┘     └───┘        EB        └───┘ │
                       ▼                                            │
                     ┌─┴─┐                                          │
                     │ Q │                                          ▼
                     └─┬─┘                                          │
                       │                                            │
                 y_hat ▼                                            │
                       │                                            │
                       ·                                            │
                    GC : ◄─────────────────────◄────────────────────┘
                       ·                 scales_hat
                       │                 means_hat
                 y_hat ▼
                       │
              ┌───┐    │
    x_hat ──◄─┤g_s├────┘
              └───┘

  文章说它是GMM(高斯混合)，但是只是预测了μ，δ，实际上是单高斯的，高斯混合应该是多个高斯模型。
  模型如下：

    def forward(self, x):
        y = self.g_a(x)
        z = self.h_a(y)
        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        gaussian_params = self.h_s(z_hat)
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        y_hat, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
        x_hat = self.g_s(y_hat)

  可以看到超先验解码输出了，方差和均值。gaussian_conditional部分输入了y，方差，均值。
  gaussian_conditional前向传播代码：

 def forward(
            self,
            inputs: Tensor,
            scales: Tensor,
            means: Optional[Tensor] = None,
            training: Optional[bool] = None,
    ) -> Tuple[Tensor, Tensor]:
        if training is None:
            training = self.training
        outputs = self.quantize(inputs, "noise" if training else "dequantize", means)
        likelihood = self._likelihood(outputs, scales, means)
        if self.use_likelihood_bound:
            likelihood = self.likelihood_lower_bound(likelihood)
        return outputs, likelihood

  对输入y进行了加均匀噪声量化，likelihood部分类似于前面的GSM，详细如下

 def _likelihood(
            self, inputs: Tensor, scales: Tensor, means: Optional[Tensor] = None
    ) -> Tensor:
        half = float(0.5)

        if means is not None:
            values = inputs - means
        else:
            values = inputs

        scales = self.lower_bound_scale(scales)

        values = torch.abs(values)
        upper = self._standardized_cumulative((half - values) / scales)
        lower = self._standardized_cumulative((-half - values) / scales)
        likelihood = upper - lower

        return likelihood

  之前的GSM没有均值部分，这次的GMM加入了均值。这段代码公式如下
$$0.5\times e^{}rfc[\frac{(x-\mu )\pm 0.5}{\sqrt{2}\sigma }]$$
  likelihood代表的区域如下：

context_prediction

来自文章

              ┌───┐    y     ┌───┐  z  ┌───┐ z_hat      z_hat ┌───┐
        x ──►─┤g_a├──►─┬──►──┤h_a├──►──┤ Q ├───►───·⋯⋯·───►──┤h_s├─┐
              └───┘    │     └───┘     └───┘        EB        └───┘ │
                       ▼                                            │
                     ┌─┴─┐                                          │
                     │ Q │                                   params ▼
                     └─┬─┘                                          │
                 y_hat ▼                  ┌─────┐                   │
                       ├──────────►───────┤  CP ├────────►──────────┤
                       │                  └─────┘                   │
                       ▼                                            ▼
                       │                                            │
                       ·                  ┌─────┐                   │
                    GC : ◄────────◄───────┤  EP ├────────◄──────────┘
                       ·     scales_hat   └─────┘
                       │      means_hat
                 y_hat ▼
                       │
              ┌───┐    │
    x_hat ──◄─┤g_s├────┘
              └───┘

EB = Entropy bottleneck
GC = Gaussian conditional
EP = Entropy parameters network#熵参数
CP = Context prediction (masked convolution)
上下文模型代码如下：

self.context_prediction = MaskedConv2d(
            M, 2 * M, kernel_size=5, padding=2, stride=1
        )
ctx_params = self.context_prediction(y_hat)

class MaskedConv2d(nn.Conv2d):
    r"""Masked 2D convolution implementation, mask future "unseen" pixels.
    Useful for building auto-regressive network components.

    Introduced in `"Conditional Image Generation with PixelCNN Decoders"
    <https://arxiv.org/abs/1606.05328>`_.

    Inherits the same arguments as a `nn.Conv2d`. Use `mask_type='A'` for the
    first layer (which also masks the "current pixel"), `mask_type='B'` for the
    following layers.
    """

    def __init__(self, *args: Any, mask_type: str = "A", **kwargs: Any):
        super().__init__(*args, **kwargs)

        if mask_type not in ("A", "B"):
            raise ValueError(f'Invalid "mask_type" value "{mask_type}"')

        self.register_buffer("mask", torch.ones_like(self.weight.data))
        _, _, h, w = self.mask.size()
        self.mask[:, :, h // 2, w // 2 + (mask_type == "B") :] = 0
        self.mask[:, :, h // 2 + 1 :] = 0

    def forward(self, x: Tensor) -> Tensor:
        # TODO(begaintj): weight assigment is not supported by torchscript
        self.weight.data *= self.mask
        return super().forward(x)

  掩膜卷积类型选择A类型，然后定义了一个和卷积核权重大小相同，內部全部为1的掩膜。通过对卷积核进行 Mask操作，掩盖了未解码点的数值，保证了自回归模型预测的当前的像素点参数仅来自于前面已经解码的点，而不取决于未解码点。通过这种Mask 掩膜，遮蔽卷积核下面和右边的权重，这种卷积核与特征图进行卷积的时候，可以可知卷积得到的结果与“未来”的数据无关。自回归模型的问题在于存在严格是时序关系，只能先得到前面的点才能得到当前点的信息，而前面的点也只能依靠更前面的点得到，即表现为在解码的时候，原始的解码方式具有并行性质，而自回归则是串行顺序。
  通过使用 self.register_buffer(name, tensor) 方法，可以将一个张量 tensor 注册为模型的缓冲区，并分配一个名称 name 给它。这意味着这个张量将与模型一起保存和加载，但不会被自动训练。
  然后获取卷积核的尺寸

self.mask[:, :, h // 2, w // 2 + (mask_type == "B") :] = 0
self.mask[:, :, h // 2 + 1 :] = 0

  上面这段代码可以看到，如果是A类型那么高度一半的下面，宽度一半的右面，包括中点，全为0。类似如下：

  如果是B类型就是类似这样，中点是1。

  A类型适用于第一层。
  forward部分：

    def forward(self, x: Tensor) -> Tensor:
        # TODO(begaintj): weight assigment is not supported by torchscript
        self.weight.data *= self.mask
        return super().forward(x)

  卷积核权重与掩膜相乘类似如下结果：

总体如下：

NLAIC

  来自End-to-End Learnt Image Compression via Non-Local Attention Optimization and Improved Context Modeling
  整体架构如下，ContextModel为上下文模型，类似前面所讲的。

              ┌───┐    X     ┌───┐  z  ┌───┐ z_hat      z_hat ┌───┐
        Y ──►─┤E_m├──►─┬──►──┤E_h├──►──┤ Q ├───►───·⋯⋯·───►──┤D_h├─┐
              └───┘    │     └───┘     └───┘        EB        └───┘ │
                       ▼                                            │
                     ┌─┴─┐                                          │
                     │ Q │                                          ▼
                     └─┬─┘                                          │
                       │                                            │P
                 X_hat ▼                                            │
                       │                                            │
                       ·                                            │
                    GC : ◄─────────────────────◄────────────────────┘
                       ·                 scales_hat
                       │                 means_hat
                 X_hat ▼
                       │
              ┌───┐    │
    Y_hat ──◄─┤D_m├────┘
              └───┘

fake, xp1, xp2, xq1, x3 = image_comp(batch_x, 1)# y^,x^_likelihood,z^_likelihood,x^,scales_hat And means_hat
xp3, _ = context(xq1, x3)#x^ ,scales_hat And means_hat  xp3:likelihoods

  上面这两个一个是加了上下文的，一个是没有上下文的，作者为了消融实验。

NLN

  非局部网络，架构如下，其理论来自本作者另一篇Non-local Attention Optimized Deep Image Compression

  基于卷积神经网络的端到端图像压缩框架是目前研究者们常用的一个网络框架，它相比其他网络框架具有稀疏连接和参数共享等特点，这有效的减少了网络的参数，网络训练检测时间也相应的减少，所以被研究者们所青睐。但是，卷积神经网络也存在局限性，卷积神经网络的感受野有限，卷积层关注和建立的联系都在局部图像范围中，也不能通过无限增加网络层数来提高感受野。
  而在图像压缩中，相邻模块的图像信息对整张图的压缩起到重要的作用，为了将图像的非局部相关性考虑进图像压缩中，在文章中，研究者设计了一 种非局部相关性模块，以获取图像之间的非局部相关信息，有利于提高图像的整体压缩效果。
  非局部模块的常见基础模型如下：

  首先我们了解一下计算机视觉领域注意力机制的开篇之作non-local operations。用于捕获长距离依赖，即如何建立图像上两个有一定距离的像素之间的联系，如何建立视频里两帧的联系，如何建立一段话中不同词的联系等。
$$Y_i=\frac{1}{C(X)}\sum _{\forall j}f(X_i,X_j)g(X_j),$$
  其中x表示输入信号(图片，序列，视频等，也可能是它们的features),y表示输出信号，其size和x相同。f(xi,xj)用来计算i和所有可能关联的位置j之间成对的关系，这个关系可以是比如i和j的位置距离越远，f值越小，表示j位置对i影响越小，g(xj)用于计算输入信号在j位置的特征值。C(x)是归一化参数。提出的non-local operations通过计算任意两个位置之间的交互直接捕捉远程依赖，non-local operations不用局限于相邻点，其相当于构造了一个和特征图谱尺寸一样大的卷积核, 从而可以维持更多信息。
  其中g由于是一元输出，比较简单，我可以采用1x1卷积，或时空中的1×1×1卷积，代表线性嵌入，${\mathrm{W}}_g$代表是要学习的权重矩阵
其形式为：$$\mathrm{g}({\mathbf{x}}_{\mathbf{j}})={\mathrm{W}}_g{\mathbf{x}}_{\mathbf{j}},$$
  对于f，前面我们说过其实就是计算两个位置的相关性，那么第一个非常自然的函数是Gaussian：$$\mathrm{f}({\mathrm{x}}_{\mathrm{i}},{\mathrm{x}}_{\mathrm{j}})={\mathrm{e}}^{{\mathrm{x}}_{\mathrm{i}}^{\mathrm{T}}{\mathrm{x}}_{\mathrm{j}}}$$
  这里${\mathbf{x}}_{\mathrm{i}}^{\mathrm{T}}{\mathbf{x}}_{\mathrm{j}}$代表点击相似性，然后通过指数映射，放大差异。标准化因子设置为$\mathrm{C}(\mathbf{x})={\sum }_{\forall \mathrm{j}}\text{f}({\mathbf{x}}_{\mathrm{i}},{\mathbf{x}}_{\mathrm{j}})$
Embedded Gaussian：
  嵌入式高斯$$\mathrm{f}({\mathrm{x}}_{\mathrm{i}},{\mathrm{x}}_{\mathrm{j}})={\mathrm{e}}^{\theta ({\mathrm{x}}_{\mathrm{i}}{)}^{\mathrm{T}}\varphi ({\mathrm{x}}_{\mathrm{j}})}$$
  高斯函数的一个简单扩展是计算嵌入空间中的相似性,在嵌入空间中计算高斯距离。$\theta \left({\mathbf{x}}_i\right)=W_{\theta }{\mathbf{x}}_i$ 和 $\varphi \left({\mathbf{x}}_j\right)=W_{\varphi }{\mathbf{x}}_j$ 是两个嵌入件。标准化因子同样$C(\mathbf{x})={\sum }_{\forall \mathrm{j}}f({\mathbf{x}}_{\mathrm{i}},{\mathbf{x}}_{\mathrm{j}})$，前面两个的标准化因子都是同样的，我们如果把C(x)考虑进去,那么$\frac{1}{\mathcal{C}(\mathbf{x})}f\left({\mathbf{x}}_i,{\mathbf{x}}_j\right)$其实就是softmax形式。完整考虑是：$$\mathbf{y}=\text{softmax}({\mathbf{x}}^T{W}_{\theta }^T{W}_{\varphi }\mathbf{x})g(\mathbf{x})$$
Dot product.
  f被定义为点积相似性,其中C(x)=N，像素个数。Dot product和Embedded Gaussian版本之间的主要区别是softmax的存在，它起着激活函数的作用。
$$\mathrm{f}({\mathrm{x}}_{\mathrm{i}},{\mathrm{x}}_{\mathrm{j}})=\theta ({\mathrm{x}}_{\mathrm{i}}{)}^{\mathrm{T}}\varphi ({\mathrm{x}}_{\mathrm{j}}).$$
Concatenation.
  这里C(x)=N
$$f\left({\mathbf{x}}_i,{\mathbf{x}}_j\right)=\mathrm{ReLU}\left({\mathbf{w}}_f^T\left[\theta \left({\mathbf{x}}_i\right),\varphi \left({\mathbf{x}}_j\right)\right]\right)$$

上面这几个就是非局部操作的算子，我们利用算子就可以构建非局部模块(如上图)。
$${\mathbf{z}}_{\mathrm{i}}={\mathrm{W}}_{\mathrm{z}}{\mathbf{y}}_{\mathrm{i}}+{\mathbf{x}}_{\mathrm{i}},$$   ${\mathbf{y}}_{\mathrm{i}}$前面我们已经给出了$+{\mathbf{x}}_{\mathrm{i}}$代表了残差连接。上面的做法的好处是可以随意嵌入到任何一个预训练好的网络中，因为只要设置${\mathrm{W}}_z$初始化为0，那么就没有任何影响，然后在迁移学习中学习新的权重。这样就不会因为引入了新的模块而导致预训练权重无法使用。
  这个就是我刚才提出的本文章作者非局部网络的理论基础。

文章NLN代码：

class Non_local_Block(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(Non_local_Block, self).__init__()
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.g = nn.Conv2d(self.in_channel, self.out_channel, 1, 1, 0)
        self.theta = nn.Conv2d(self.in_channel, self.out_channel, 1, 1, 0)
        self.phi = nn.Conv2d(self.in_channel, self.out_channel, 1, 1, 0)
        self.W = nn.Conv2d(self.out_channel, self.in_channel, 1, 1, 0)
        nn.init.constant_(self.W.weight, 0)
        nn.init.constant_(self.W.bias, 0)

    def forward(self, x):
        # x_size: (b c h w)

        batch_size = x.size(0)
        g_x = self.g(x).view(batch_size, self.out_channel, -1)
        g_x = g_x.permute(0, 2, 1)
        theta_x = self.theta(x).view(batch_size, self.out_channel, -1)
        theta_x = theta_x.permute(0, 2, 1)
        phi_x = self.phi(x).view(batch_size, self.out_channel, -1)

        f1 = torch.matmul(theta_x, phi_x)
        f_div_C = f.softmax(f1, dim=-1)
        y = torch.matmul(f_div_C, g_x)
        y = y.permute(0, 2, 1).contiguous()
        y = y.view(batch_size, self.out_channel, *x.size()[2:])
        W_y = self.W(y)
        z = W_y+x

        return z

整体流程如下：

NLAM

  NLAM 代表非局部注意力模块,主分支由三个 ResBlock 组成。 mask 分支将非本地模块与 ResBlock 结合起来，用于生成注意力 mask。具有通道和空间的注意力机制。
代码如下：

class ResBlock(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride, padding):
        super(ResBlock, self).__init__()
        self.in_ch = int(in_channel)
        self.out_ch = int(out_channel)
        self.k = int(kernel_size)
        self.stride = int(stride)
        self.padding = int(padding)

        self.conv1 = nn.Conv2d(self.in_ch, self.out_ch,
                               self.k, self.stride, self.padding)
        self.conv2 = nn.Conv2d(self.in_ch, self.out_ch,
                               self.k, self.stride, self.padding)

    def forward(self, x):
        x1 = self.conv2(f.relu(self.conv1(x)))
        out = x+x1
        return out
self.trunk5 = nn.Sequential(ResBlock(self.M, self.M, 3, 1, 1), ResBlock(self.M, self.M, 3, 1, 1),
                                    ResBlock(self.M, self.M, 3, 1, 1))
self.mask2 = nn.Sequential(Non_local_Block(self.M, self.M // 2), ResBlock(self.M, self.M, 3, 1, 1),
                                   ResBlock(self.M, self.M, 3, 1, 1),
                                   ResBlock(self.M, self.M, 3, 1, 1), nn.Conv2d(self.M, self.M, 1, 1, 0))                                  
x6 = self.trunk5(x5)*f.sigmoid(self.mask2(x5)) + x5#NLAM

流程如下：

3D Masked Convolution

  我前面提到了Minnen等人提出的通过每个特征通道的2D的掩膜卷积来提取自回归信息。本文作者在他的NLAIC中提出了使用3D掩膜卷积，比2D增加了一个通道维度，它可以联合利用空间和跨通道相关性。传统 2D PixelCNN 需要搜索所有先前通道以利用条件概率。本文提出的 3D 掩码卷积隐式捕获相邻通道之间的相关性。基于 3D CNN 的方法通过在整个空间通道空间中强制执行相同的卷积核，显着减少了条件上下文建模的网络参数。通过自回归方式利用空间通道邻居的附加上下文，可以获得条件高斯分布。

  看上面的图，画××的方块是我们当前要预测的像素，黄色的方块是上一个通道的像素，绿色方块是它当前通道垂直方向相邻的像素，蓝色的方块是它当前通道水平相邻的像素，利用这些已知的部分来预测当前像素。空白的方块是未处理的，它和当前像素被0掩膜覆盖。
  代码如下：

class MaskConv3d(nn.Conv3d):
    def __init__(self, mask_type, in_ch, out_ch, kernel_size, stride, padding):
        super(MaskConv3d, self).__init__(in_ch, out_ch,
                                         kernel_size, stride, padding, bias=True)

        self.mask_type = mask_type
        ch_out, ch_in, k, k, k = self.weight.size()
        mask = torch.zeros(ch_out, ch_in, k, k, k)
        central_id = k*k*k//2+1
        current_id = 1
        if mask_type == 'A':
            for i in range(k):
                for j in range(k):
                    for t in range(k):
                        if current_id < central_id:
                            mask[:, :, i, j, t] = 1
                        else:
                            mask[:, :, i, j, t] = 0
                        current_id = current_id + 1
        if mask_type == 'B':
            for i in range(k):
                for j in range(k):
                    for t in range(k):
                        if current_id <= central_id:
                            mask[:, :, i, j, t] = 1
                        else:
                            mask[:, :, i, j, t] = 0
                        current_id = current_id + 1
        self.register_buffer('mask', mask)

    def forward(self, x):
        self.weight.data *= self.mask
        return super(MaskConv3d, self).forward(x)

  获取卷积核权重形状，然后设置一个形状相同但是全为0的掩膜，获取卷积核中心位置的序号，从1开始计数，由于卷积核是一个三维的立方体，所以中心位置的id等于卷积核体积的一半加1，之前部分为1，后面为0，并且A类型屏蔽了中心。

GAACNN

  联合图注意力和非对称卷积神经网络（GAACNN)。来自Joint Graph Attention and Asymmetric Convolutional Neural Network for Deep Image Compression

  前人的一些图像压缩方法通过平等对待所有像素来采用不同的视觉特征，而忽略了局部关键特征的影响。仍然需要探索更有效的注意力机制来将比特分配给不同的图像区域。虽然用于图像压缩的自注意力机制显示出性能提升，但这些机制只探索了空间级别或通道级别的冗余，而忽略了跨通道和位置之间的依赖性。这些注意力机制的另一个限制是嵌入在不同层的特征图中的注意力信息无法相互交互，因为这些机制只处理当前时刻提取的特征图的注意力关系。虽然抽象表示中的长程依赖可以通过堆叠 CNN来建模，但这是一个次优解决方案，因为重复堆叠的卷积运算使得优化图像压缩变得困难。
  为了克服上述困难作者提出了联合图注意力和非对称卷积神经网络GAACNN
所提出的GAACNN不仅考虑了高级抽象表示中的远程依赖性，而且突出了潜在表示中局部特征的影响。
  GAACNN由图注意力网络（GAT），非对称卷积神经网络（ACNN），自注意机制主要构成。

ACNN(非对称卷积神经网络)

  一般来说，定制的CNN通过平等地处理所有像素来聚合不同的特征，以提高图像压缩的性能，但忽略了局部特征的影响，过高估计了非关键特征的作用。本质上，与方形核卷积相比，水平和垂直方向上的1×3和3×1的1D非对称卷积核可以关注方形核的局部区域。也就是说，这两个1D非对称卷积用于增强局部显著特征，以表示重要的细节。因此，提出了由方形核卷积和1D非对称卷积组成的ACNN模块，以低成本地通过增强局部关键特征的作用来强调图像细节。
  局部关键特征通过水平和垂直方向上的两个1D非对称卷积得到增强。之后，将增强的局部关键特征聚合到一个方形核卷积中，以实现全局和局部信息的结合。
  同时，通过将方形卷积核分解为两个较小的1D卷积核，可以减少模型参数的数量，从而加速深度压缩网络的训练。
  ACNN架构如图：

  非对称块由三个卷积层和一个激活函数组成，即Conv1、Conv2、Conv3和ReLU，Conv1和Conv3是1D非对称卷积，分别包含一个3×1和一个1×3的卷积核。Conv2是一个3×3的方形卷积核，它在残差学习过程中会受到Conv1和Conv3的影响，以优化提取的局部关键特征，并增强图像压缩网络对局部细节的表达能力，ReLU用于增强提取特征的非线性变换。
  ACNN的特征提取过程定义为：$$\begin{array}{rl}{O}_i^C& =C_1(O_{i-1}^R)+C_2(O_{i-1}^R)+C_3(O_{i-1}^R),i\ge 1,\\ F& =\sum _{i=1}^3{O}_i^C,\end{array}$$
  其中，R代表了ReLU，C1,C2,C3分别代表了Conv1、Conv2、Conv3，$O_{i-1}^R$代表了上一层经过ReLU激活后卷积特征，$O_i^C$是ACNN中第i层的卷积特征，就是Asymmetric Block。 F表示ACNN的输出。
  代码：

class ACBlock(nn.Module):#Asymmetric Block
    def __init__(self, in_channels, out_channels):
        super(ACBlock, self).__init__()
        self.conv1x3 = nn.Sequential(nn.Conv2d(in_channels, out_channels, (1, 3), 1, (0, 1)))
        self.conv3x1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, (3, 1), 1, (1, 0)))
        self.conv3x3 = nn.Sequential(nn.Conv2d(in_channels, out_channels, (3, 3), 1, (1, 1)))

    def forward(self, x):
        conv3x1 = self.conv3x1(x)
        conv1x3 = self.conv1x3(x)
        conv3x3 = self.conv3x3(x)
        return conv3x1 + conv1x3 + conv3x3
for idx in range(1, 4):
	self.__setattr__(f"blk{idx + 1}", nn.Sequential(nn.ReLU(inplace=True), ACBlock(N2 // 2, N2 // 2)))
for idx in range(1, 4):#ACNN Moudle
            tmp = self.__getattr__(f"blk{idx + 1}")(tmp)
            high1 = high1 + tmp

Attention Module

  本文提出的提出的注意力网络将通道注意力和空间注意力结合了起来。
  首先了解一下空间注意力机制和通道注意力机制：
  通道注意力机制
  1.SEnet

  SENet 分为压缩和激励两部分，其中压缩部分的目的是对全局空间信息进行压缩，然后在通道维度进行特征学习，形成各个通对道的重要性，最后通过激励部分对各个通道进行分配不同权重的。
  $F_{s\mathbf{q}}(\cdot )$全局平均池化，将每个通道的二维特征（H*W）压缩为1个实数，将特征图从 [h, w, c] ==> [1,1,c]
  $F_{ex}(\cdot ,w)$给每个特征通道生成一个权重值，论文中通过两个全连接层构建通道间的相关性，激活分别是relu,sigmoid，输出的权重值数目和输入特征图的通道数相同。[1,1,c] ==> [1,1,c]
  $F_{scale}(\cdot )$将前面得到的归一化权重加权到每个通道的特征上。论文中使用的是乘法，逐通道乘以权重系数。[h,w,c]×[1,1,c] ==> [h,w,c]
  混合注意力机制
  1.CBAM
  CBAM注意力机制是由通道注意力机制（channel）和空间注意力机制（spatial）组成。CBAM的总体流程图如下。输入特征图先经过通道注意力机制，将通道权重和输入特征图相乘后再送入空间注意力机制，将归一化后的空间权重和空间注意力机制的输入特征图相乘，得到最终加权后的特征图。

空间部分：
  对通道注意力机制的输出特征图进行空间域的处理。首先，对输入特征图在通道维度下做最大池化和平均池化，将池化后的两张特征图在通道维度堆叠。然后，使用 7×7 （或3×3、1×1）大小的卷积核融合通道信息，特征图的shape从 [b,2,h,w] 变成 [b,1,h,w]。最后，将卷积后的结果经过 sigmoid 函数对特征图的空间权重归一化，再将输入特征图和权重相乘。

本文的通道注意力机制使用了多个堆叠的通道注意力残差块，通道注意力机制通过残差块学习通道之间的依赖关系，动态调整哪些通道相对重要。

  空间注意力机制如下，空间注意机制将特征映射分解为两个一维特征编码过程，在水平方向和垂直方向上保留位置信息，增强显著区域的表征，定位目标位置。

  集成输出如下，其中$T_j^c\mathrm{和}{T}_j^s$第j个通道注意模块和第j个空间注意模块的输出。α和β是平衡通道注意和空间注意的两个超参数。
$$T_j=\alpha \cdot T_j^c+\beta \cdot T_j^s,j\ge 1,$$

  将前一个注意模块的输出传递给当前注意模块，参与当前注意图的计算，有助于相邻注意模块之间的信息共享。
  信息交互后的输出定义如下,${\overline{T}}_{j+1}$表示第(j+1)个注意模块之后的信息交互输出，将第j个注意模块的特征与第(j+1)个注意模块的特征进行加权得到。γ和η是两个可学习的参数，用于动态交互输出。注意，当γ = 0 和 η = 1时，注意模块之间不存在共享信息。

$${\overline{T}}_{j+1}=\gamma \cdot T_j+\eta \cdot T_{j+1},j\ge 1,$$
  通道部分代码如下：

class CALayer(nn.Module):
    def __init__(self, channel, reduction=1):
        super(CALayer, self).__init__()
        # global average pooling: feature --> point
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        # feature channel downscale and upscale --> channel weight
        self.conv_du = nn.Sequential(
                nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
                nn.Sigmoid()
        )

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.conv_du(y)
        return x * y
class RCAB(nn.Module):
    def __init__(
        self, conv, n_feat, kernel_size, reduction=1,
        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):

        super(RCAB, self).__init__()
        modules_body = []
        for i in range(2):
            modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
            if bn: modules_body.append(nn.BatchNorm2d(n_feat))
            if i == 0: modules_body.append(act)
        modules_body.append(CALayer(n_feat, reduction))
        self.body = nn.Sequential(*modules_body)
        self.res_scale = res_scale

    def forward(self, x):
        res = self.body(x)
        res = self.body(x).mul(self.res_scale)
        res += x
        return res

  空间部分代码如下：

class CoordAtt(nn.Module):
    def __init__(self,inp,oup,reduction=32):
        super(CoordAtt,self).__init__()
        self.pool_h=nn.AdaptiveAvgPool2d((None,1))#垂直方向上的输出的高度将固定为1
        self.pool_w=nn.AdaptiveAvgPool2d((1,None))

        mip=max(8,inp//reduction)

        self.conv1=nn.Conv2d(inp,mip,1,1,0)
        self.bn1=nn.BatchNorm2d(mip)
        self.act=h_swish()

        self.conv_h=nn.Conv2d(mip,oup,1,1,0)
        self.conv_w=nn.Conv2d(mip,oup,1,1,0)

    def forward(self,x):
        identity=x
        n,c,h,w=x.size()#n：批处理大小（Batch Size）或样本数量。
                        #c：通道数（Channels），即输入张量中特征通道的数量。
                        #h：高度（Height），即输入张量的高度维度的大小。
                        #w：宽度（Width），即输入张量的宽度维度的大小。
        x_h=self.pool_h(x)#垂直方向上的输出的高度将固定为1
        x_w=self.pool_w(x).permute(0,1,3,2)

        y=torch.cat([x_h,x_w],dim=2)
        y=self.conv1(y)
        y=self.bn1(y)
        y=self.act(y)

        x_h,x_w=torch.split(y,[h,w],dim=2)
        x_w=x_w.permute(0,1,3,2)

        a_h=self.conv_h(x_h).sigmoid()
        a_w=self.conv_w(x_w).sigmoid()

        out=identity*a_w*a_h

        return out

  最终整体的Attention Module由残差通道注意力机制和空间注意力机制得到的结果相加。self.attention1(tmp_1)是空间注意力机制，encode1_channel是残差堆叠的通道注意力机制。tmp_1就是完整的一个Attention Module块。

tmp_1=self.attention1(tmp_1)+encode1_channel

GAT(图注意力网络)

  参考文章图卷积网络
  
  
  
  

LEARNING ACCURATE ENTROPY MODEL WITH GLOBAL REFERENCE FOR IMAGE COMPRESSION

  来自文章

  作者认为：上下文自适应模型 [Minnen et al, 2018a, Lee et al, 2019] 结合了相邻符号的预测，以避免存储额外的位。虽然这些方法提高了熵模型的准确性，但它们无法在压缩过程中使用全局上下文信息，从而导致性能不佳。
  由于观察到全局空间冗余仍然存在于潜在变量中，如图 1 所示。受此启发，我们建议在整个潜在变量中建立全局相关性。受到最近基于参考的超分辨率（SR）方法的启发[Zheng et al, 2018, Yang et al, 2020]，我们通过合并参考组件为熵模型赋予了全局视野。
  在我们提出的方法中，全局参考模块搜索已解码的潜在变量，以找到与目标潜在变量相关的潜在变量。然后，相关潜在特征图与局部上下文和超先验相结合，生成更准确的熵估计。
  不仅考虑相关目标与目标之间的相似性，而且还考虑用于测量潜在特征分布中的高阶静态的置信度得分。置信度得分的引入增强了熵模型的鲁棒性，特别是对于具有噪声背景的图像。
  我们发现图像压缩中广泛使用的广义除法归一化（GDN）存在均值平移问题。由于 GDN 密度根据定义为零均值，因此需要去除均值来拟合密度 [Ballé et al, 2016b]。因此，我们提出了 GDN 的改进版本，称为 GSDN（广义减法和除法标准化）来克服这个困难。
  主体的Encoder和Decoder我们沿用了之前的方法。对于概率模型，我们加入了全局相关性搜索的模块，并且设计了更为精细的概率模型网络。本方法的概率模型结合了局部上下文预测模块（Local Context）、全局搜索参考模块（Global Reference）、超先验概率模块（Hyperprior）。

encoder:

  

hyper_encoder:

  

hyper_decoder:

  

全局参考模块

  相似度矩阵：
  在解码目标潜在变量时，我们使用相邻潜在变量（左侧和顶部）作为基础来计算目标潜在变量与其先前潜在变量之间的相似性。特别是，潜伏被展开成补丁，然后被掩蔽，表示为 q ∈ [H × W, k × k × C] （其中 H、W 、k、C 分别对应于高度、宽度、展开内核大小和通道）。我们使用余弦相似度计算整个掩模斑块的相似度矩阵 r ∈ [H × W, H × W ]$$r_{i,j}=⟨\frac{q_i}{\Vert q_i\Vert },\frac{q_j}{\Vert q_j\Vert }⟩$$
  代码部分：

class SearchTransfer(nn.Module):
    def __init__(self, C, k=3, split=1): #384，3
        super().__init__()
        # modified
        # Mask Type 1
        mask = torch.ones((C // split, k, k))# 384 ， 3×3
        mask[:, k // 2, k // 2:] = 0#中心点右下方
        mask[:, k // 2 + 1:, :] = 0#中心点下方
        mask_unfold = F.unfold(mask.unsqueeze(0), kernel_size=(k, k), padding=0)#只提取特征，不进行积运算。
        self.mask_unfold = torch.nn.Parameter(mask_unfold, requires_grad=False)#1，3456，1
        self.k = k
        self.split = split

    def forward(self, y_hat, y_prob):
        k = self.k

        # Search 输入是[batchsize, channel, h,w]
        # 输出是[batchsize,channel* kH * kW, L] 其中kH是kernel的高，kW是kernel宽。 L则是这个高kH宽kW的kernel能在H*W区域按照指定stride滑动的次数。
        #每一行带表了，卷积提取的元素特征，不进行积运算

        unfold = F.unfold(y_hat, kernel_size=(k, k), padding=k//2) * self.mask_unfold #(1,3456,1024)×(1,3456,1)
        unfold = F.normalize(unfold, dim=1)  # [N, C*k*k, H*W]
        unfold_T = unfold.permute(0, 2, 1)  # [N, H*W, C*k*k]

        R = torch.bmm(unfold_T, unfold)

  R为相似度矩阵，和它的转置相乘，相当于求内积， F.unfold函数在做卷积的滑动窗口。但是只有卷的部分，没有积的部分，举例如下。

x = torch.arange(0, 1*1*5*5).float()
x = x.view(1,1,5,5)
print(x)
x1 = F.unfold(x, kernel_size=3, dilation=1, stride=1, padding=0)
print(x1.shape)
print(x1)

结果如下：

tensor([[[[ 0.,  1.,  2.,  3.,  4.],
          [ 5.,  6.,  7.,  8.,  9.],
          [10., 11., 12., 13., 14.],
          [15., 16., 17., 18., 19.],
          [20., 21., 22., 23., 24.]]]])
torch.Size([1, 9, 9])
tensor([[[ 0.,  1.,  2.,  5.,  6.,  7., 10., 11., 12.],
         [ 1.,  2.,  3.,  6.,  7.,  8., 11., 12., 13.],
         [ 2.,  3.,  4.,  7.,  8.,  9., 12., 13., 14.],
         [ 5.,  6.,  7., 10., 11., 12., 15., 16., 17.],
         [ 6.,  7.,  8., 11., 12., 13., 16., 17., 18.],
         [ 7.,  8.,  9., 12., 13., 14., 17., 18., 19.],
         [10., 11., 12., 15., 16., 17., 20., 21., 22.],
         [11., 12., 13., 16., 17., 18., 21., 22., 23.],
         [12., 13., 14., 17., 18., 19., 22., 23., 24.]]])

  经过F.unfold函数的矩阵可以理解为用卷积提取，每一行都带表了一个卷积提取的元素特征，不进行积运算。

继续看：

 unfold = F.normalize(unfold, dim=1)  # [N, C*k*k, H*W]
        unfold_T = unfold.permute(0, 2, 1)  # [N, H*W, C*k*k]

        R = torch.bmm(unfold_T, unfold) #[N, H*W, H*W] #对应batchsize相乘

        # Refer all when trainning, modified
        if(self.training):
            R = torch.triu(R, diagonal=1) + torch.tril(R, diagonal=-1)#上三角矩阵和下三角矩阵
        else:
            R = torch.triu(R, diagonal=1)

        R_star, R_star_arg = torch.max(R, dim=1)

 R_star, R_star_arg = torch.max(R, dim=1)

  这里返回每一行的最大值以及在每一行的位置，位置从0开始索引。之前说过q矩阵每一行都是一个卷积核滑动提取的特征而且每一行都是经过掩膜卷积掩膜卷积掩蔽过的。

 U_unfold = self.bis(unfold_prob, 2, R_star_arg)

    def bis(self, input, dim, index):#y_hat_unfold, 2, R_star_arg
        # batch index select
        # input: [N, ?, ?, ...]
        # dim: scalar > 0
        # index: [N, idx]
        views = [input.size(0)] + [1 if i!=dim else -1 for i in range(1, len(input.size()))]#1，1，-1生成的列表的形状与输入张量 input 的维度数减去 1 相同，且在维度 dim 上的元素值为 -1，其他位置的元素值为 1。
        # 这样的列表主要用于构造 views 和 expanse 变量，以便在后续的操作中创建适当形状的索引张量。
        expanse = list(input.size())
        expanse[0] = -1
        expanse[dim] = -1
        index = index.view(views).expand(expanse)

  上面代码的意思是将index索引矩阵扩展为和input一样的形状，并且索引矩阵每一行都是相同的，行数和input相同。

return torch.gather(input, dim, index)

返回index所对应位置的input

  相似度S和置信度U

unfold_prob = F.unfold(y_prob, kernel_size=(1, 1), padding=0)
        U_unfold = self.bis(unfold_prob, 2, R_star_arg) # [N, 1, H*W]

        n, c, h, w = y_hat.shape
        S = R_star.view(n, 1, h, w)
        U = F.fold(U_unfold, output_size=(h,w), kernel_size=(1,1), padding=0)
 if not self.training:
            S[:,:,0,0], U[:,:,0,0], ref_unfold[:,:,0], R_star_arg[:,0]  = 1e-8, 1e-8, 0., -1

        S = torch.clamp(S, min=1e-8, max=1.0)
        U = torch.clamp(U, min=1e-8, max=1.0)

        return S, U, ref_unfold, R_star_arg#相似度S和置信度U

  decoder:

  代码注意部分
  在算数编码解码部分计算value = ((offset + 1) * total - 1) // range的时候值会溢出。

 def read(self, freqs):
    if not isinstance(freqs, CheckedFrequencyTable):
      freqs = CheckedFrequencyTable(freqs)
    
    # Translate from coding range scale to frequency table scale
    total = freqs.get_total()
    if total > self.maximum_total:
      raise ValueError("Cannot decode symbol because total is too large")
    range = self.high - self.low + 1
    offset = self.code - self.low
    offset = np.int64(offset)
    range = np.int64(range)
    value = ((offset + 1) * total - 1) // range

  需要把range和offset加上np.int64。