python - 可以在 Python 中更快地进行这种 inpaint 插值吗?
问题描述
根据这篇论文Garcia et . 人。(2012 年)。我尝试将此代码(inpaintn.m)移植到 Python 中,如下所示,
import numpy as np
from scipy.ndimage import distance_transform_edt
from scipy.fft import idctn, dctn
from tqdm import tqdm
def fill_nd(data, invalid=None):
if invalid is None: invalid = np.isnan(data)
ind = distance_transform_edt(invalid, return_distances=False, return_indices=True)
return data[tuple(ind)]
def InitialGuess(y, I):
z = fill_nd(y)
s0 = 3
return z, s0
def idctnn(y):
return idctn(y, norm='ortho')
def dctnn(y):
return dctn(y, norm='ortho')
def inpaint(xx, y0=[], n=100, m=2, verbose=False):
x = xx.copy() #as it changes x itself, so copying it to another variable.
sizx = np.shape(x)
d = np.ndim(x)
Lambda = np.zeros(sizx, dtype='float')
for i in range(0, d):
siz0 = np.ones(d, dtype='int')
siz0[i] = sizx[i]
Lambda = Lambda + np.cos(np.pi * np.reshape(np.arange(1, sizx[i] + 0.1) - 1, siz0) / sizx[i])
Lambda = 2 * (d - Lambda)
# Initial condition
W = np.isfinite(x)
if len(y0) == len(x):
y = y0
s0 = 3 # note: s = 10 ^ s0
else:
if np.any(~W):
if verbose: print('Initial Guess as Nearest Neighbors')
y, s0 = InitialGuess(x, np.isfinite(x).astype('bool'))
else:
y = x
s0 = 3
# return x
x[~W] = 0.
# Smoothness parameters: from high to negligible
s = np.logspace(s0, -6, n)
RF = 2. # Relaxation Factor
Lambda = Lambda ** m
if verbose: print('Inpainting .......')
for i in tqdm(range(n)):
Gamma = 1. / (1 + s[i] * Lambda)
y = RF * idctnn(Gamma * dctnn((W * (x - y)) + y)) + (1 - RF) * y
y[W] = x[W]
return y
该代码运行良好,但我一直在努力寻找使该代码运行得更快的方法,尤其是因为我的数据集很大。使用这种插值的优点是我可以输入整个 3D 数据集(带有时间和网格坐标)来填充缺失值,而不是为每个时间坐标都这样做。
这是一个使用 python 的示例数据集
import numpy as np
#A 3D dataset with dimensions (time, latitude, longitude)
X = np.random.randn(1000,180,360)
# Randomly choosing indices to insert 64800 NaN values (say).
#NaNs can also be present as blocks in the data, not randomly dispersed as below.
index_nan = np.random.choice(X.size, 64800, replace=False)
#Inserting NaNs.
X.ravel()[index_nan] = np.nan
我尝试了一些方法,但都没有成功,
- 使用 Numba
jit 装饰器使它变慢,即使使用parallel/fastmath/vectorize,nopython=True
.
- 使用 Cython
我尝试排版这些函数中使用的所有变量,但它仍然比本机 python 实现慢。而且,在我的机器上编译 Cython 代码很麻烦。
- 使用 Numpy 向量化
我已经用函数替换了离散余弦变换函数及其逆scipy
函数,但我似乎无法想出将内部 for 循环矢量化以使其快速的方法,以及它是否可能。我已经尝试分析我的代码,瓶颈似乎是在离散余弦变换中使用scipy
. 还有其他瓶颈,但对我来说没有意义。我还附上了一张用于分析的图片。
如果有可行的方法来加速这段代码,那真的会有很大帮助。我在 Python 方面并不是很先进,但是我可以从中学到很多东西,尤其是我的问题的可行性。
解决方案
该算法适用于一个相当大的数组(不适合 CPU 缓存),部分解释了为什么它有点慢。此外,众所周知,DCT/IDCT 是昂贵的操作。话虽如此,您可以使用 Numba 的 JIT 和workers=-1
scipy 函数选项来并行化算法。此外,您可以通过就地工作来避免创建许多昂贵的临时阵列。这是未经测试的结果代码:
# In-place computation
def idctnn(y):
return idctn(y, norm='ortho', workers=-1, overwrite_x=True)
# In-place computation
def dctnn(y):
return dctn(y, norm='ortho', workers=-1, overwrite_x=True)
# In-place computation (writes in `Transformed`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64)', parallel=True)
def ComputeGammaTransform(Transformed, Lambda, sVal):
for i in nb.prange(Transformed.shape[0]):
for j in range(Transformed.shape[1]):
for k in range(Transformed.shape[2]):
Transformed[i, j, k] /= (1. + sVal * Lambda[i, j, k])
# Out-of-place computation (writes in `out`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64[:,:,::1], boolean[:,:,::1])', parallel=True)
def ComputeDctInput(out, x, y, W):
for i in nb.prange(out.shape[0]):
for j in range(out.shape[1]):
for k in range(out.shape[2]):
out[i, j, k] = W[i, j, k] * (x[i, j, k] - y[i, j, k]) + y[i, j, k]
# In-place computation (writes in `y`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64)', parallel=True)
def ComputeDctOutput(dctResult, y, RF):
for i in nb.prange(y.shape[0]):
for j in range(y.shape[1]):
for k in range(y.shape[2]):
y[i, j, k] = RF * dctResult[i, j, k] + (1.0 - RF) * y[i, j, k]
def ComputeSteps(Lambda, x, y, W, s, RF):
dctData = np.empty(Lambda.shape, dtype=Lambda.dtype)
for i in tqdm(range(s.shape[0])):
ComputeDctInput(dctData, x, y, W)
dctnn(dctData)
ComputeGammaTransform(dctData, Lambda, s[i])
idctnn(dctData)
ComputeDctOutput(dctData, y, RF)
这段代码在我的机器上快了 5 倍。您可以使用简单精度而不是双精度来进一步加快速度。这使得最终代码比我机器上的原始代码快 7.5 倍。
我也许可以通过基于 GPU 的计算进一步加快代码速度。困难的部分是在 Python 中找到支持正交归一化的 DCT/IDCT 的 GPU 实现。
推荐阅读
- node.js - 过了一会儿我得到了这个错误:错误:没有通道可以分配
- javascript - 动态类绑定 svg 和 vue 不起作用
- pcap - 我在哪里可以找到 Authentication-Information-Request 的所有 avp Pcap?
- python - 如何找到对应于另一个 3D 对象的旋转和平移矩阵
- sql-server - 查找以 GUID 作为名称的表
- plsql - PL/SQL 触发器在更新或插入后更新同一个表
- maven - Cucumber 报告在每个 Jenkins 构建上触发一个错误
- powershell - 使用 PowerShell 删除 CSV 文件中的空行
- c# - 从开始和结束位置拆分字符串
- javascript - Firebase自动生成密钥时,是否可以更新孩子的数据?