Reimplementation of GradNotSetToNonePattern from Torchtidy (#92)

Adding rules to check for `set_to_none` parameter for `zero_grad()`.

By setting set_to_none=True, we can gain speedup

Author: shivam096

tests/fixtures/performance/checker/zerograd.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+x = torch.ones((100, 100))
+model = nn.Sequential()
+optimizer = torch.optim.Adam(model.parameters())
+
+# This should raise flags
+optimizer.zero_grad(set_to_none=False)
+model.zero_grad(set_to_none=False)
+
+# This should not raise flags 
+optimizer.zero_grad()
+model.zero_grad()
+
+

tests/fixtures/performance/checker/zerograd.txt

@@ -0,0 +1,2 @@
+9:1 TOR402 Detected gradient set to zero instead of None. Please add 'set_to_none=True' when calling zero_grad().
+10:1 TOR402 Detected gradient set to zero instead of None. Please add 'set_to_none=True' when calling zero_grad().

torchfix/torchfix.py

@@ -21,6 +21,7 @@
     TorchVisionDeprecatedPretrainedVisitor,
     TorchVisionDeprecatedToTensorVisitor,
     TorchVisionSingletonImportVisitor,
+    TorchGradNotSetToNonePatternVisitor,
 )
 
 __version__ = "0.7.0"
@@ -43,6 +44,7 @@
     TorchVisionDeprecatedPretrainedVisitor,
     TorchVisionDeprecatedToTensorVisitor,
     TorchVisionSingletonImportVisitor,
+    TorchGradNotSetToNonePatternVisitor,
 ]
 
 

torchfix/visitors/__init__.py

@@ -8,7 +8,10 @@
     TorchRequireGradVisitor,
 )
 from .nonpublic import TorchNonPublicAliasVisitor
-from .performance import TorchSynchronizedDataLoaderVisitor
+from .performance import (
+    TorchSynchronizedDataLoaderVisitor,
+    TorchGradNotSetToNonePatternVisitor,
+)
 from .security import TorchUnsafeLoadVisitor
 from .vision import (
     TorchVisionDeprecatedPretrainedVisitor,
@@ -30,4 +33,5 @@
     "TorchVisionDeprecatedPretrainedVisitor",
     "TorchVisionDeprecatedToTensorVisitor",
     "TorchVisionSingletonImportVisitor",
+    "TorchGradNotSetToNonePatternVisitor",
 ]

torchfix/visitors/performance/__init__.py

@@ -32,3 +32,36 @@ def visit_Call(self, node):
                     error_code=self.ERRORS[0].error_code,
                     message=self.ERRORS[0].message(),
                 )
+
+
+class TorchGradNotSetToNonePatternVisitor(TorchVisitor):
+    """
+    Reimplementation of GradNotSetToNonePattern from
+    https://github.com/pytorch/pytorch/blob/main/torch/profiler/_pattern_matcher.py
+    """
+
+    ERRORS = [
+        TorchError(
+            "TOR402",
+            (
+                "Detected gradient set to zero instead of None. "
+                "Please add 'set_to_none=True' when calling zero_grad()."
+            ),
+        )
+    ]
+
+    def visit_Call(self, node):
+        qualified_name = self.get_qualified_name_for_call(node)
+
+        if qualified_name and qualified_name.endswith("zero_grad"):
+
+            set_to_none_arg = self.get_specific_arg(node, "set_to_none", 0)
+
+            # hasattr check to handle mypy error
+            if set_to_none_arg and hasattr(set_to_none_arg.value, "value"):
+                if set_to_none_arg.value.value == "False":
+                    self.add_violation(
+                        node,
+                        error_code=self.ERRORS[0].error_code,
+                        message=self.ERRORS[0].message(),
+                    )