首页 > 解决方案 > 在张量流中为张量层解包位时性能缓慢

问题描述

我正在处理通过 WebSocket 连接与星际争霸 2 客户端连接的数据,以从正在进行的游戏中获取图像数据。在某些情况下,图像数据可以设置为每像素 1 位的格式。发生这种情况时,我需要从响应中的每个字节“解包”位(1 字节 => 8 位)。这是在下面的代码中完成的:

function unpackbits(uint8data) {
  const results = new Uint8Array(8 * uint8data.length)
  let byte
  let offset
  for (let i = 0; i < uint8data.length; i++) {
    byte = uint8data[i]
    offset = (8 * i)
    results[offset + 7] = ((byte & (1 << 0)) >> 0)
    results[offset + 6] = ((byte & (1 << 1)) >> 1)
    results[offset + 5] = ((byte & (1 << 2)) >> 2)
    results[offset + 4] = ((byte & (1 << 3)) >> 3)
    results[offset + 3] = ((byte & (1 << 4)) >> 4)
    results[offset + 2] = ((byte & (1 << 5)) >> 5)
    results[offset + 1] = ((byte & (1 << 6)) >> 6)
    results[offset + 0] = ((byte & (1 << 7)) >> 7)
  }
  return results
}

这被输入到一个张量中,如下所示:

 static unpack_layer(plane) {
    //Return a correctly shaped tensor given the feature layer bytes.//

    const size = point.Point.build(plane.getSize()) // { x, y }

    if (plane.getBitsPerPixel() === 1) {
      data = unpackbits(data)
      if (data.length !== (size.x * size.y)) {
        // This could happen if the correct length isn't a multiple of 8, leading
        // to some padding bits at the end of the string which are incorrectly
        // interpreted as data.
        data = data.slice(0, size.x * size.y)
      }
    }

    data = tf.tensor(data, [size.y, size.x], 'int32')
    return data
}

在我的一项测试中,这段代码运行了 1900 次,执行时间为 0.0737 秒。

这是非常缓慢的。

为了比较,python 中的等效功能需要 0.0209 秒才能运行 1900 次。python代码如下所示:

  def unpack_layer(plane):
    """Return a correctly shaped numpy array given the feature layer bytes."""

    size = point.Point.build(plane.size) # {x, y }
    data = np.frombuffer(plane.data, dtype=Feature.dtypes[plane.bits_per_pixel])

    if plane.bits_per_pixel == 1:
      data = np.unpackbits(data)
      if data.shape[0] != size.x *  size.y:
        # This could happen if the correct length isn't a multiple of 8, leading
        # to some padding bits at the end of the string which are incorrectly
        # interpreted as data.
        data = data[:size.x * size.y]
    return data.reshape(size.y, size.x)

简而言之,javascript 版本大约是 python 版本的 4 倍。

我将查看 numpy unpackbits 文档,因为这似乎比我自己的方法更有效 -

但是,我想知道是否有人对如何更好地优化我自己的 unpackbits 函数或更好的方法让 TensorFlow 为我做这件事有任何想法?

标签: javascriptbitmapbit-manipulationtensorflow.js

解决方案


不确定这是否有帮助,但我很自责,因为我对 tensorflow 中的位运算符的需求感到困惑,以便根据原始问题将字节流转换为位流。简单地使用整数除法和模数也可以解决问题!

简而言之,示例算法就是这样。给定 [92] 的字节流...

  • 除以 16 并取模,得到 2 个字节,分别为 [5] 和 [12]。
  • 将这些结果交织成张量 [5, 12]。
  • 取其中的每一个值,然后除以 4 并取模,得到 [1, 3] 和 [1, 0]。
  • 将这些结果交织成张量 [1, 1, 3, 0]。
  • 除以 2 并取模,得到 [ 0, 0, 1, 0 ] 和 [ 1, 1, 1, 0 ]。
  • 交织成 [ 0, 1, 0, 1, 1, 1, 0, 0 ],它是 92 的二进制。

以下是同一算法的两个版本。一个在 tensorflow 中,一个在纯 JavaScript 中。

function tfDaC( stream ) {
  
  const stream8bit = tf.tensor( stream, undefined, 'int32' );
  
  console.time('in-tf');
  const stream4bitHi = tf.div(stream8bit, tf.scalar(16, 'int32' ));
  const stream4bitLo = tf.mod(stream8bit, tf.scalar(16, 'int32' ));
  const stream4bit = tf.stack([stream4bitHi, stream4bitLo],1).flatten();

  const stream2bitHi = tf.div( stream4bit, tf.scalar(4, 'int32' ));
  const stream2bitLo = tf.mod(stream4bit, tf.scalar(4, 'int32' ));
  const stream2bit = tf.stack([stream2bitHi, stream2bitLo],1).flatten();

  const stream1bitHi = tf.div(stream2bit, tf.scalar(2, 'int32' ));
  const stream1bitLo = tf.mod(stream2bit, tf.scalar(2, 'int32' ));
  const stream1bit = tf.stack([stream1bitHi, stream1bitLo],1).flatten().toBool();
  console.timeEnd('in-tf');
  
  return stream1bit.dataSync().buffer;
}


function jsDaC( stream ) {

  let result = new ArrayBuffer( stream.byteLength * 8 );

  let buffer32 = new Uint32Array( result );  // Pointer to every 4 bytes!  
  for ( let i = 0; i < stream.byteLength; i++ ) {
    let byte = stream[ i ];
    buffer32[ (i * 2) |0 ] = ( byte / 16) |0;
    buffer32[ (i * 2 + 1) |0 ] = ( byte % 16 ) |0;
  }
  
  let buffer16 = new Uint16Array( result );  // Pointer to every 2 bytes!  
  for ( let i = 0; i < buffer32.length; i++ ) {
    let byte = buffer32[ i ];
    buffer16[ (i * 2) |0 ] = ( byte / 4) |0;
    buffer16[ (i * 2 + 1) |0 ] = ( byte % 4 ) |0;
  }
  
  let buffer8 = new Uint8Array( result );  // Pointer to every 4 bytes!  
  for ( let i = 0; i < buffer16.length; i++ ) {
    let byte = buffer16[ i ];
    buffer8[ (i * 2) |0 ] = ( byte / 2 ) |0;
    buffer8[ (i * 2 + 1) |0 ] = ( byte % 2 ) |0;
  }

  return result;
}

console.log( 'Generating array of 1M bytes' );
let buffer = new ArrayBuffer( 1000000 );
let testArray = new Uint8Array( buffer );
for ( let i = 0; i < testArray.length; i++ ) {
  testArray[ i ] = Math.floor( 256 * Math.random() );
}

let result;

console.log( 'Begin tensorflow divide & conquer test with 1M bytes.' );
console.time( 'tf' );
result = tfDaC( testArray );
console.timeEnd( 'tf' );
console.log( `End tensorflow test with 1M bytes resulting in array of ${result.byteLength} bytes` );

console.log( 'Begin javascript divide & conquer test with 1M bytes.' );
console.time( 'js' );
result = jsDaC( testArray );
console.timeEnd( 'js' );
console.log( `End javascript test with 1M bytes resulting in array of ${result.byteLength} bytes` );
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.0.1/dist/tf.min.js"></script>

我的工作站上的 tensorflow 性能很糟糕。我不得不将字节流减少到 1M 字节,因为我的 GPU 在我之前的 10M 字节流测试级别引发内存错误。即便如此,只有 1M 字节,少数测试从 1236 毫秒到 1414 毫秒不等。不知道为什么这么慢。可能是将数字强制转换为 int32,这可能会增加很多开销,因为我的理解是 GPU 通常主要用于浮点运算。并且将数据编组到 GPU 上和从 GPU 上输出也需要一些时间。也许值得尝试将此函数转换为仅浮点函数而不是 int32 ...?!也许是一个糟糕的 tensorflow.js 版本......?!有兴趣了解它如何在您的 NodeJS 配置中运行...

另一方面,1M 字节的 javascript 版本从 30 毫秒到 42 毫秒不等,几乎比 GPU 快 2 个数量级(!)。但是,当将这些结果外推到 10M 字节时,该算法仍然比所有其他以前的算法慢...

所以不确定这是否有帮助。它可能只是帮助消除 tensorflow 作为一种选择,尽管尝试浮点数而不是 int32 可能仍然值得,但我不是很有希望......


推荐阅读