SageMaker TF 2.3 分布式训练

Question

使用SageMaker v2.29.2和Tensorflow v2.3.2，我正在尝试实现分布式训练，如下面的博客文章所述：

https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-customize-training-script-tf.html#model-parallel-customize-training-script-tf-23

然而，我在导入smdistributed脚本时遇到了困难。

这是我的代码：

import tensorflow as tf
import smdistributed.modelparallel.tensorflow as smp

错误：

Traceback (most recent call last):
  File "temp.py", line 2, in <module>
    import smdistributed.modelparallel.tensorflow as smp
ModuleNotFoundError: No module named 'smdistributed'

我错过了什么？

- Shlomi Schwartz

---

您可以尝试使用Tensorflow 2.3.1或2.4.1，并告诉我们是否解决了您的问题？ - user11530462

Answer 1

smdistributed仅在SageMaker容器上可用。它支持特定的TensorFlow版本，您必须添加：

distribution={'smdistributed': {
            'dataparallel': {
                'enabled': True
            }
        }}

在估算器代码中使其启用

- Shlomi Schwartz