コード例 #1
0
ファイル: test_loader.py プロジェクト: stjordanis/TensorRT
 def test_unmark_shape_outputs(self, modifiable_reshape_network):
     builder, network, parser = modify_network_outputs(
         modifiable_reshape_network,
         outputs=constants.MARK_ALL,
         exclude_outputs=["reduce_prod_out_gs_2"])
     with builder, network, parser:
         assert network.num_outputs == 1
コード例 #2
0
 def test_exclude_outputs_with_mark_layerwise(self, modifiable_network):
     builder, network, parser = modify_network_outputs(
         modifiable_network, outputs=constants.MARK_ALL, exclude_outputs=["identity_out_2"]
     )
     with builder, network, parser:
         assert network.num_outputs == 1
         assert network.get_output(0).name == "identity_out_0"
コード例 #3
0
ファイル: test_loader.py プロジェクト: stjordanis/TensorRT
 def test_mark_shape_outputs(self, modifiable_reshape_network):
     builder, network, parser = modify_network_outputs(
         modifiable_reshape_network,
         outputs=["output", "reduce_prod_out_gs_2"])
     with builder, network, parser:
         assert network.num_outputs == 2
         assert network.get_output(0).name == "reduce_prod_out_gs_2"
         assert network.get_output(0).is_shape_tensor
コード例 #4
0
ファイル: test_loader.py プロジェクト: stjordanis/TensorRT
 def test_mark_custom_outputs(self, modifiable_network):
     builder, network, parser = modify_network_outputs(
         modifiable_network, outputs=["identity_out_0"])
     with builder, network, parser:
         assert network.num_outputs == 1
         assert network.get_output(0).name == "identity_out_0"