optimize cast op

x2paddle/decoder/tf_decoder.py

@@ -60,6 +60,15 @@ class TFGraphNode(GraphNode):
             raise Exception("Dtype[{}] not in dtype_map".format(dtype))
         return self.dtype_map[dtype]
 
+    @property
+    def raw_dtype(self):
+        keys = ['dtype', 'Tidx', 'T', 'DstT']
+        for k in keys:
+            dtype = self.layer.attr[k].type
+            if dtype > 0:
+                break
+        return dtype
+
     @property
     def value(self):
         assert self.layer_type == "Const", "Only Const node has value."
@@ -120,6 +129,7 @@ class TFGraph(Graph):
         # tensorflow graph optimize
         self._remove_isolated_node()
         self._remove_identity_node()
+        self._remove_cast_node()
 
     def get_node(self, node_name, copy=False):
         items = node_name.strip().split(':')
@@ -190,6 +200,27 @@ class TFGraph(Graph):
                 idx = self.output_nodes.index(node_name)
                 self.output_nodes[idx] = input_node.layer_name
 
+    def _remove_cast_node(self):
+        cast_node = list()
+        for node_name, node in self.node_map.items():
+            if node.layer_type == "Cast":
+                input = self.get_node(node.inputs[0])
+                if input.layer_type != "Placeholder" or len(input.outputs) != 1:
+                    continue
+                cast_node.append(node_name)
+
+        for node_name in cast_node:
+            node = self.get_node(node_name)
+            input_node = self.get_node(node.inputs[0])
+            input_node.layer.attr["dtype"].type = node.raw_dtype
+            self.remove_node(node_name)
+
+            self.identity_map[node_name] = input_node.layer_name
+
+            if node_name in self.output_nodes:
+                idx = self.output_nodes.index(node_name)
+                self.output_nodes[idx] = input_node.layer_name
+
     def data_format_propagation(self, node):
         current_node = self.node_map[node.layer_name]
         current_node = node.tf_data_format