(self)
| 23 | |
| 24 | class TestJointGraph(unittest.TestCase): |
| 25 | def test_joint_graph(self) -> None: |
| 26 | class Module(torch.nn.Module): |
| 27 | def __init__(self): |
| 28 | super().__init__() |
| 29 | self.linear = torch.nn.Linear(3, 3) |
| 30 | self.linear_no_train = torch.nn.Linear(3, 3) |
| 31 | for param in self.linear_no_train.parameters(): |
| 32 | param.requires_grad = False |
| 33 | self.loss = torch.nn.CrossEntropyLoss() |
| 34 | |
| 35 | def forward(self, x, y): |
| 36 | return self.loss(self.linear_no_train(self.linear(x)).softmax(dim=0), y) |
| 37 | |
| 38 | m = Module() |
| 39 | example_inputs = (torch.ones(3), torch.tensor([1.0, 0.0, 0.0])) |
| 40 | m(*example_inputs) |
| 41 | ep = _export(m, example_inputs, pre_dispatch=True) |
| 42 | joint_ep = _export_forward_backward(ep) |
| 43 | edge = to_edge(joint_ep) |
| 44 | |
| 45 | output_node = edge.exported_program().graph.output_node() |
| 46 | |
| 47 | orig_outputs = len(output_node.args[0]) |
| 48 | |
| 49 | et = edge.to_executorch() |
| 50 | |
| 51 | weight_output_specs = [ |
| 52 | spec |
| 53 | for spec in et.exported_program().graph_signature.output_specs |
| 54 | if spec.kind == OutputKind.TOKEN |
| 55 | ] |
| 56 | |
| 57 | output_node = et.exported_program().graph.output_node() |
| 58 | |
| 59 | weight_outputs = len(output_node.args[0]) |
| 60 | |
| 61 | # make sure 2 new outputs are added to both the node and the spec |
| 62 | self.assertEqual(len(weight_output_specs), 2) # linear layer weight and bias |
| 63 | self.assertEqual( |
| 64 | weight_outputs - orig_outputs, 2 |
| 65 | ) # linear layer weight and bias |
| 66 | |
| 67 | # assert that the weight and bias have proper data_buffer_idx and allocation_info |
| 68 | self.assertEqual( |
| 69 | et.executorch_program.execution_plan[0].values[0].val.data_buffer_idx, |
| 70 | 1, |
| 71 | ) |
| 72 | self.assertEqual( |
| 73 | et.executorch_program.execution_plan[0].values[1].val.data_buffer_idx, |
| 74 | 2, |
| 75 | ) |
| 76 | self.assertEqual( |
| 77 | et.executorch_program.execution_plan[0] |
| 78 | .values[0] |
| 79 | .val.allocation_info.memory_offset_low, |
| 80 | 96, |
| 81 | ) |
| 82 | self.assertEqual( |
nothing calls this directly
no test coverage detected