| @@ -79,23 +79,24 @@ def test_decode_layer_03(): | |||||
| d = Data() | d = Data() | ||||
| d.add_node_type('Dummy 1', 100) | d.add_node_type('Dummy 1', 100) | ||||
| d.add_node_type('Dummy 2', 100) | d.add_node_type('Dummy 2', 100) | ||||
| d.add_relation_type('Dummy Relation 1', 0, 1, | |||||
| fam = d.add_relation_family('Dummy 1-Dummy 2', 0, 1, True) | |||||
| fam.add_relation_type('Dummy Relation 1', | |||||
| torch.rand((100, 100), dtype=torch.float32).round().to_sparse()) | torch.rand((100, 100), dtype=torch.float32).round().to_sparse()) | ||||
| prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) | prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) | ||||
| in_layer = OneHotInputLayer(d) | in_layer = OneHotInputLayer(d) | ||||
| d_layer = DecagonLayer(in_layer.output_dim, 32, d) | d_layer = DecagonLayer(in_layer.output_dim, 32, d) | ||||
| dec_layer = DecodeLayer(input_dim=d_layer.output_dim, data=prep_d, keep_prob=1., | |||||
| decoder_class={(0, 1): DEDICOMDecoder}, activation=lambda x: x) | |||||
| dec_layer = DecodeLayer(input_dim=d_layer.output_dim, data=prep_d, | |||||
| keep_prob=1., activation=lambda x: x) | |||||
| seq = torch.nn.Sequential(in_layer, d_layer, dec_layer) | seq = torch.nn.Sequential(in_layer, d_layer, dec_layer) | ||||
| pred_adj_matrices = seq(None) | |||||
| assert isinstance(pred_adj_matrices, dict) | |||||
| assert len(pred_adj_matrices) == 2 | |||||
| assert isinstance(pred_adj_matrices[0, 1], list) | |||||
| assert isinstance(pred_adj_matrices[1, 0], list) | |||||
| assert len(pred_adj_matrices[0, 1]) == 1 | |||||
| assert len(pred_adj_matrices[1, 0]) == 1 | |||||
| pred = seq(None) | |||||
| assert isinstance(pred, Predictions) | |||||
| assert len(pred.relation_families) == 1 | |||||
| assert isinstance(pred.relation_families[0], RelationFamilyPredictions) | |||||
| assert isinstance(pred.relation_families[0].relation_types, list) | |||||
| assert len(pred.relation_families[0].relation_types) == 1 | |||||
| assert isinstance(pred.relation_families[0].relation_types[0], RelationPredictions) | |||||
| def test_decode_layer_04(): | def test_decode_layer_04(): | ||||