()
| 26 | assert a.net.transformer.position_embeddings.weight.shape[-1] == 128 |
| 27 | |
| 28 | def test_save_and_load(): |
| 29 | args = NestedModel.get_args( |
| 30 | num_layers=2, |
| 31 | bert_args={"hidden_size": 128, "num_layers": 3, 'num_types': 2} |
| 32 | ) |
| 33 | args.mode = 'inference' |
| 34 | args.save = './checkpoints/test_nested_model' |
| 35 | args.tokenizer_type = 'fake' |
| 36 | a = NestedModel(args=args).cuda() |
| 37 | |
| 38 | from sat.training.model_io import save_checkpoint |
| 39 | |
| 40 | save_checkpoint(1, a, None, None, args) |
| 41 | |
| 42 | assert os.path.exists( |
| 43 | os.path.join(args.save, str(1), 'mp_rank_00_model_states.pt')) |
| 44 | |
| 45 | b, args = NestedModel.from_pretrained(args.save) |
| 46 | |
| 47 | assert b.net.transformer.position_embeddings.weight.shape[-1] == 128 |
| 48 | |
| 49 | # compare the weights equal between a and b |
| 50 | for a_p, b_p in zip(a.parameters(), b.parameters()): |
| 51 | assert torch.allclose(a_p, b_p) |
| 52 | |
| 53 | def test_load(): |
| 54 | trained_dir = './checkpoints/test_train_nested/MyModel-05-11-01-35' |
no test coverage detected