内存缩减Tensorflow TPU v2/v3 bfloat16

Question

我的型号太大了，不能用普通的v2 TPU设备来获得一个batch>；64。在troubleshooting网站上提到，即将推出的tensorflow版本将支持bfloat16。新支持的tf版本1.9-1.12现在是否能够使用bfloat16？如果是，是否有一组有限的优化器可以使用？我没有找到关于这个的任何进一步的文档，但是在tensor2tensor模型中看到了bfloat16的用法，所以我想肯定有一种方法。在

此外，我读到了TPU v3 supports bigger models as well但是模型需要最小的更改，但是我没有找到任何需要更改的文档。在

我已经在使用Adafactor并试图减少我的层，如果你有任何进一步的减少技巧，那也太好了。我使用图片矩阵和单词向量（目前是float32）作为输入。在

Answer 1

您可以将bfloat16与tpu一起使用。主要有两件事要做：

1. 将输入转换为输入管道中的bfloat16
2. 在bfloat16范围内环绕您的网络，并将输出转换为F32以进行进一步计算。在

下面是一个代码片段，说明了必要的更改：

def input_fn():

  def dataset_parser(self, value):
    """Parse an ImageNet record from a serialized string Tensor."""
    image = self.image_preprocessing_fn(
        image_bytes=image_bytes,
        is_training=self.is_training,
    )

    if self.use_bfloat16:
      image = tf.cast(image, tf.bfloat16)

    return image, label


def resnet_model_fn(features, labels, mode, params):
  """The model_fn for ResNet to be used with TPUEstimator."""

  # This nested function allows us to avoid duplicating the logic which
  # builds the network, for different values of  precision.
  def build_network():
    network = resnet_model.resnet_v1(
        resnet_depth=FLAGS.resnet_depth,
        num_classes=LABEL_CLASSES,
        data_format=FLAGS.data_format)
    return network(
        inputs=features, is_training=(mode == tf.estimator.ModeKeys.TRAIN))

  if FLAGS.precision == 'bfloat16':
    with bfloat16.bfloat16_scope():
      logits = build_network()
    logits = tf.cast(logits, tf.float32)
  elif FLAGS.precision == 'float32':
    logits = build_network()

您还可以看到this TPU model中说明的第二个条件。在