使用 tensorflow 和 gpu 计算像素答案

【问题标题】：counting pixels using tensorflow and gpu使用 tensorflow 和 gpu 计算像素
【发布时间】：2018-09-17 07:34:04
【问题描述】：

我有大小为 (N, 256, 256) 的蒙版图像，其中 N 是介于 1000-10000 之间的值。
每个像素都有一个介于 0-2 之间的整数值（0 只是背景）。
不幸的是，掩码图像未编码为 (N,256,256,2)
我有几千个这样的口罩。我的目标是找到最快的方法来计算每个标签（1 和 2）的每帧像素。
使用 numpy 在一张大约 6000 帧的蒙版图像上运行需要

np.sum(ma==1,axis=(1,2))
np.sum(ma==2,axis=(1,2))

如果我使用单进程，我预计运行整个数据需要几个小时，如果我使用多处理 (CPU)，可能不到一个小时。我很好奇如果我使用 GPU 是否可以让它更快。实现对轴上的张量求和的部分似乎很容易，但我找不到如何在 tensorflow 上实现 ma==1 部分。
我考虑过首先将输入输入到编码形状 (N,256,256,2) 并传递给张量占位符，但意识到制作具有该形状的数组需要比上面更长的时间。或者，有没有更好的方法来使用 tensorflow 在这个掩码数据上实现像素计数？

【问题讨论】：

我不太相信它会产生巨大的影响，因为这样的任务需要线性时间，并且将图像加载到 GPU 内存中也需要线性时间。

标签： python numpy tensorflow

【解决方案1】：

想想后台发生了什么

在您的原始实施中，以下步骤大致执行了两次：

从内存中加载整个数组，证明某个值是否等于所需值
将结果写回内存（临时数组与输入数组一样大，假设为 np.uint8）
将整个数组加载到内存中并总结结果
将结果写回内存

应该清楚的是，无论是否并行化，这都是一个非常不理想的实现。我无法以纯矢量化的 numpy 方式做得更好，但是有可用的工具（Numba、Cython），您可以在其中以更直接和并行化的方式实现此任务。

示例

import numpy as np
import numba as nb
import time

#Create some data
N=10000
images=np.random.randint(0, high=3, size=(N,256,256), dtype=np.uint8)

def sum_orig(ma):
  A=np.sum(ma==1,axis=(1,2))
  B=np.sum(ma==2,axis=(1,2))
  return A,B

@nb.njit(fastmath=True,parallel=True)
def sum_mod(ma):
  A=np.zeros(ma.shape[0],dtype=np.uint32)
  B=np.zeros(ma.shape[0],dtype=np.uint32)

  #parallel loop
  for i in nb.prange(ma.shape[0]):
    AT=0
    BT=0
    for j in range(ma.shape[1]):
      for k in range(ma.shape[2]):
        if (ma[i,j,k]==1):
          AT+=1
        if (ma[i,j,k]==2):
          BT+=1

    A[i]=AT
    B[i]=BT

  return A,B

#Warm up
#The funtion is compiled at the first call
[A,B]=sum_mod(images)
t1=time.time()
[A,B]=sum_mod(images)
print(time.time()-t1)
t1=time.time()
[A_,B_]=sum_orig(images)
print(time.time()-t1)

#check if it works correctly
print(np.allclose(A,A_))
print(np.allclose(B,B_))

性能

improved_version: 0.06s
original_version: 2.07s
speedup: 33x

【讨论】：