使用tensorflow和gpu计算像素

Question

我有大小（N，256,256）的掩模图像，其中N是1000-10000之间的值 每个像素的整数值介于0-2之间（0只是背景） 不幸的是，掩模图像不被编码为（N，256,256,2）
我有几千个这样的面具。我的目标是找到最快的方法来计算每个标签（1和2）的每帧像素数 使用numpy在大约6000帧的一个掩模图像上运行＆lt; 2秒。

  

np.sum（毫安== 1，轴=（1,2））
     np.sum（毫安== 2，轴=（1,2））

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

Answer 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