sagemaker_debugger_builtin_rules_vanishing_gradient_config.py

1import smdebug_rulesconfig as rules
2from sagemaker.tensorflow import TensorFlow
3
4# Initialize the built-in rule configuration
5# This example uses the VanishingGradient rule
6vanishing_gradient_rule = {
7    "rule_to_invoke": "VanishingGradient",
8    "rule_parameters": {
9        "threshold": "0.0001"
10    }
11}
12
13# Alternatively, using the helper function from smdebug_rulesconfig
14# to get the full configuration dictionary required by SageMaker
15built_in_rules = [
16    rules.vanishing_gradient(),
17    rules.check_vanishing_gradient(),
18    rules.loss_not_decreasing()
19]
20
21# Use the rules in a SageMaker Estimator
22estimator = TensorFlow(
23    role="SageMakerRole",
24    image_uri="<your-training-image-uri>",
25    instance_count=1,
26    instance_type="ml.p3.2xlarge",
27    rules=[
28        rules.Rule.built_in(rules.vanishing_gradient()),
29        rules.Rule.built_in(rules.loss_not_decreasing())
30    ]
31)