sagemaker_debugger_vanishing_gradient_rule_with_tensorflow_estimator.py

1import smdebug_rulesconfig as rulesconf
2from sagemaker.tensorflow import TensorFlow
3
4# Initialize the built-in rule configuration
5vanishing_gradient_rule = rulesconf.vanishing_gradient()
6
7# Define the SageMaker Estimator and include the rule
8estimator = TensorFlow(
9    role='SageMakerRole',
10    image_uri='<training-container-image-uri>',
11    instance_count=1,
12    instance_type='ml.p3.2xlarge',
13    rules=[
14        {
15            "rule_name": vanishing_gradient_rule["rule_name"],
16            "source_uri": vanishing_gradient_rule["source_uri"],
17            "instance_type": "ml.t3.medium",
18            "rule_parameters": vanishing_gradient_rule["rule_parameters"]
19        }
20    ]
21)
22
23estimator.fit(inputs='s3://my-bucket/training-data')