sagemaker_debugger_vanishing_gradient_rule_with_smdebug_rulesconfig.py

1import smdebug_rulesconfig as rulesconfig
2from sagemaker.tensorflow import TensorFlow
3
4# Initialize the built-in rule configuration
5# This example uses the VanishingGradient rule
6check_vanishing_gradient = rulesconfig.vanishing_gradient()
7
8# Use the configuration in a SageMaker Estimator
9estimator = TensorFlow(
10    role='sagemaker_role',
11    instance_count=1,
12    instance_type='ml.p3.2xlarge',
13    entry_point='train.py',
14    framework_version='2.3.1',
15    py_version='py37',
16    # Add the rule to the estimator
17    rules=[
18        rulesconfig.Rule.sagemaker(check_vanishing_gradient)
19    ]
20)
21
22estimator.fit(wait=False)