SageMaker PyTorchModel 传递自定义变量

Question

通过类使用 SageMaker 部署模型时PyTorchModel，是否可以传递自定义环境变量或 kwargs？

我希望能够通过自定义参数切换服务代码的功能，而不是需要编写多个serve.py来处理不同的训练模型导出方法。

model = PyTorchModel(name='my_model',
                     model_data=estimator.model_data,
                     role=role,
                     framework_version='1.0.0',
                     entry_point='serve.py',
                     source_dir='src',
                     sagemaker_session=sess,
                     predictor_cls=ImagePredictor,
                     <custom_argument?>
                    )

Answer 1

您是否尝试过使用env您的 中的参数PyTorchModel？（参见https://sagemaker.readthedocs.io/en/stable/model.html#sagemaker.model.Model）

model = PyTorchModel(name='my_model',
                     model_data=estimator.model_data,
                     role=role,
                     framework_version='1.0.0',
                     entry_point='serve.py',
                     source_dir='src',
                     sagemaker_session=sess,
                     predictor_cls=ImagePredictor,
                     env={'ENV_VALUE': 'val'}
                    )