内存减少 Tensorflow TPU v2/v3 bfloat16
Memory reduction Tensorflow TPU v2/v3 bfloat16
我的模型太大,无法使用普通 v2 TPU 设备获得大于 64 的批次。在 troubleshooting 网站上提到即将推出的 tensorflow 版本将支持 bfloat16。新支持的 tf 版本 1.9-1.12 现在是否能够使用 bfloat16?如果是,我可以使用的优化器是否有限?我没有找到任何关于此的进一步文档,但在 tensor2tensor 模型中看到了 bfloat16 的用法,所以我想一定有办法。
此外,我读到 TPU v3 supports bigger models as well 但该模型需要进行最少的更改,但我没有找到任何需要更改的文档。
我已经在使用 Adafactor 并尝试减少我的层数,如果您有任何进一步的减少技巧,那也很好。我正在使用图片矩阵和词向量(目前为 float32)作为输入。
您可以将 bfloat16
与 TPU 结合使用。主要有两件事要做:
- 将输入转换为输入管道中的 bfloat16
- 将您的网络包围在 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中说明的第二个条件。
我的模型太大,无法使用普通 v2 TPU 设备获得大于 64 的批次。在 troubleshooting 网站上提到即将推出的 tensorflow 版本将支持 bfloat16。新支持的 tf 版本 1.9-1.12 现在是否能够使用 bfloat16?如果是,我可以使用的优化器是否有限?我没有找到任何关于此的进一步文档,但在 tensor2tensor 模型中看到了 bfloat16 的用法,所以我想一定有办法。
此外,我读到 TPU v3 supports bigger models as well 但该模型需要进行最少的更改,但我没有找到任何需要更改的文档。
我已经在使用 Adafactor 并尝试减少我的层数,如果您有任何进一步的减少技巧,那也很好。我正在使用图片矩阵和词向量(目前为 float32)作为输入。
您可以将 bfloat16
与 TPU 结合使用。主要有两件事要做:
- 将输入转换为输入管道中的 bfloat16
- 将您的网络包围在 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中说明的第二个条件。