Positional Encoding高频位置编码

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)

# 类似 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

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()

# 观察 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)))

# 接下来是使用 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)

# 再创建一个神经网络, 这次每条输入的维度是 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()

# 观察 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)))