Positional Encoding高频位置编码

下图范例来自CVPR 2020 Tutorial, 左边是真实图片, 中间是不用位置编码的神经网络拟合结果, 右边采用了位置编码更加清晰。

本项目将复现类似的结果。        

import paddlefrom matplotlib import pyplot as plt 
import numpy as np 
from tqdm import tqdm 
from PIL import Image

准备工作

很少有更简单的数据准备工作了——随便找一张图片即可。

# 随便找一张图片, 放进项目的文件夹里并用 Image 库打开img = Image.open('work/example1.jpg')
img = np.array(img)print(img.shape)
plt.figure(figsize=(5,7))
plt.imshow(img)

(960, 720, 3)

<matplotlib.image.AxesImage at 0x7fb48d122290>

<Figure size 360x504 with 1 Axes>

# 类似 NeRF 的想法, 用一个神经网络表示图片：给一个像素坐标 (x,y), 预测该像素点的 RGB 值。# 神经网络简单地用若干 Linear 层 和 ReLU 激活函数堆叠而成。class NeRF2D(paddle.nn.Layer):
    def __init__(self, input_size: int = 2, layers: int = 5, hidden_size: int = 256):
        super().__init__()        assert layers >= 2, 'MLP should have at least 2 layers.'
        input_size :int = input_size # 2d 输入: 像素坐标 (x,y)
        output_size:int = 3          # 3d 输出: 坐标的对应RGB值 (R,G,B)
    
        # 保存一些参数
        self.input_size = input_size 
        self.output_size = output_size 
        self.layers = layers 

        # 用一个 Layerlist 存放所有层
        self.mlps = paddle.nn.LayerList()
        mlp_dim = [input_size] + [hidden_size] * (layers - 1) + [output_size]        for layer in range(layers):            # 全连接层
            self.mlps.append(
                paddle.nn.Linear(mlp_dim[layer], mlp_dim[layer+1])
            )            # ReLU层
            self.mlps.append(paddle.nn.ReLU())    def forward(self, x):
        # 让 x 通过所有层
        for layer in range(self.layers * 2): 
            x = self.mlps[layer](x)        return x

不使用 Positional Encoding

用 $F(x,y)$F(x,y) 表示神经网络函数 ($(x,y)$(x,y) 是输入), 则损失函数为神经网络的预测 $F(x,y)$F(x,y) 与真实图片的像素值 $RGB(x,y)$RGB(x,y) 的差距 (L2 损失)。

$$L=\frac{1}{N}\sum _{\mathrm{a}\mathrm{l}\mathrm{l}\mathrm{p}\mathrm{i}\mathrm{x}\mathrm{e}\mathrm{l}\mathrm{s}(x,y)}{(F(x,y)-RGB(x,y))}^2$$L=N1all pixels (x,y)∑(F(x,y)−RGB(x,y))2

w , h = img.shape[0], img.shape[1]
pi = np.pi 
# 将 w*h 个像素坐标 (x,y) 作为输入# 这里不妨将图片看成 [-pi/2, pi/2] x [-pi/2, pi/2] 的坐标系# (因为神经网络的输入和输出最好不要太大, 每个值 <= 1 比较好)inputs = paddle.zeros((w,h,2))
inputs[:,:,0] += paddle.linspace(-pi/2, pi/2, num=w).reshape((w,1)) # 纵坐标inputs[:,:,1] += paddle.linspace(-pi/2, pi/2, num=h).reshape((1,h)) # 横坐标# 输入相当于 w*h 条二维数据 (输入不是很大, 可以一次性全部计算, 不需要 batch)inputs = inputs.reshape((w*h, 2))# 作为参照的, 期望的输出是 w*h 个像素的 RGB 值, 即图片本身outputs = paddle.to_tensor(img, dtype='float32').reshape((w*h, 3))
outputs = outputs / 255. # 由于网络最后是 ReLU, 所以期望的输出应该 >= 0, 最好介于 [0,1]

# 创建一个神经网络, 用 Adam 优化器net = NeRF2D()
opt = paddle.optimizer.Adam(parameters = net.parameters())
losses = []

# 训练 2000 步 (3分钟左右)epochs = 2000for epoch in tqdm(range(len(losses)+1, len(losses)+epochs+1)):
    render = net(inputs) # (w*h, 3)
    # 将神经网络根据坐标预测的渲染 (render) 结果与原图片对比, 用 L2 损失作为损失函数
    loss = paddle.mean(paddle.square(render - outputs))
    losses.append(loss.item())    # 反向传播+梯度下降
    opt.clear_grad()
    loss.backward()
    opt.step()

100%|██████████| 2000/2000 [03:06<00:00, 10.70it/s]

# 观察 loss 曲线 和 最后一次渲染的结果, # 神经网络生成的图片有模有样, 缺点是比较模糊plt.figure(figsize=(12,7))
plt.subplot(1,2,1)
plt.semilogy(losses)
plt.subplot(1,2,2)
plt.imshow(render.reshape((w,h,3)))

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

<matplotlib.image.AxesImage at 0x7fb48ce69510>

<Figure size 864x504 with 2 Axes>

使用 Positional Encoding

使用了 Positional Encoding, 即将每个输入从 $(x,y)$(x,y) 变成

$$[x,y,\sin (x),\sin (y),\sin (2x),\sin (2y),\mathrm{.}\mathrm{.}\mathrm{.},\cos (x),\cos (y),\cos (2x),\cos (2y),\mathrm{.}\mathrm{.}\mathrm{.}]$$[x,y,sin(x),sin(y),sin(2x),sin(2y),...,cos(x),cos(y),cos(2x),cos(2y),...]

这里我们最高取到 $2^0,2^1,2^2,2^3,2^4$20,21,22,23,24 倍.

# 接下来是使用 Positional Encoding 的改进# 这时候每个输入不再是 2 维的 (x,y)# 而是多维的 (x,y,sinx,siny,cosx,cosy,sin(2x),sin(2y),cos(2x),cos(2y),...)positional_encoding = [inputs]
freqs = [1.,2.,4.,8.,16.]for freq in freqs:
    positional_encoding.append(paddle.cos(inputs * freq))
    positional_encoding.append(paddle.sin(inputs * freq))

positional_encoding = paddle.to_tensor(positional_encoding, dtype='float32') # shape = [1+2*Freqs, w*h, 2]positional_encoding = positional_encoding.transpose((1,2,0)).reshape((w*h, -1))print(positional_encoding.shape)

[691200, 22]

# 再创建一个神经网络, 这次每条输入的维度是 22net = NeRF2D(input_size = positional_encoding.shape[1])
opt = paddle.optimizer.Adam(parameters = net.parameters())
losses = []

# 训练 2000 步 (3分钟左右)epochs = 2000for epoch in tqdm(range(len(losses)+1, len(losses)+epochs+1)):
    render = net(positional_encoding) # (w*h, 3)
    # 将神经网络根据坐标预测的渲染 (render) 结果与原图片对比, 用 L2 损失作为损失函数
    loss = paddle.mean(paddle.square(render - outputs))
    losses.append(loss.item())

    opt.clear_grad()
    loss.backward()
    opt.step()

100%|██████████| 2000/2000 [03:08<00:00, 10.62it/s]

# 观察 loss 曲线 和 最后一次渲染的结果, 可以看出比不用 positional encoding 要清晰很多plt.figure(figsize=(12,7))
plt.subplot(1,2,1)
plt.semilogy(losses)
plt.subplot(1,2,2)
plt.imshow(render.reshape((w,h,3)))

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).

<matplotlib.image.AxesImage at 0x7fb48cc63450>

<Figure size 864x504 with 2 Axes>