fix pir to kwargs

PaddlePaddle · Sep 20, 2023 · 05067af · 05067af
1 parent bdaff69
commit 05067af
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 10 deletions.
diff --git a/demo/components/program_test.py → demo/components/pir_translate.py b/demo/components/program_test.py → demo/components/pir_translate.py
@@ -25,4 +25,4 @@
         model=newir_program,
         input_spec=[paddle.static.InputSpec([-1, 1, 28, 28], 'float32')],
         verbose=True,
-        is_newir=True)
+        is_pir=True)
diff --git a/visualdl/component/graph/exporter.py b/visualdl/component/graph/exporter.py
@@ -22,8 +22,8 @@
 from .utils import print_model
 
 
-def translate_graph(model, input_spec, verbose=True, is_pir=False):
-    import paddle
+def translate_graph(model, input_spec, verbose=True, **kwargs):
+    is_pir = kwargs.get('is_pir', False)
     with tempfile.TemporaryDirectory() as tmp:
         if (not is_pir):
             model._full_name = '{}[{}]'.format(model.__class__.__name__,

diff --git a/visualdl/component/graph/utils.py b/visualdl/component/graph/utils.py
@@ -30,9 +30,9 @@ def gen_var_name(ops):
         if var in var_name:
             return var_name[var]
         else:
-            if (op.is_persistable):
+            try:
                 name = op.name
-            else:
+            except ValueError:
                 name = "tmp_var_" + str(var_idx[0])
                 var_idx[0] += 1
             var_name[var] = name

diff --git a/visualdl/writer/writer.py b/visualdl/writer/writer.py
@@ -16,7 +16,6 @@
 import time
 
 import numpy as np
-from paddle.framework import in_pir_mode
 
 from visualdl.component.base_component import audio
 from visualdl.component.base_component import embedding
@@ -614,7 +613,7 @@ def add_roc_curve(self,
                       num_thresholds=num_thresholds,
                       weights=weights))
 
-    def add_graph(self, model, input_spec, verbose=False):
+    def add_graph(self, model, input_spec, verbose=False, **kwargs):
         """
         Add a model graph to vdl graph file.
         Args:
@@ -655,9 +654,8 @@ def forward(self, inputs):
                     verbose=True)
         """
         try:
-            if in_pir_mode():
-                is_pir = True
-            result = translate_graph(model, input_spec, verbose, is_pir=False)
+            is_pir = kwargs.get('is_pir', False)
+            result = translate_graph(model, input_spec, verbose, is_pir=is_pir)
         except Exception as e:
             print("Failed to save model graph, error: {}".format(e))
             raise e