| 95 | EXPECTED_TOTALS[epoch - 1]) |
| 96 | |
| 97 | def test_single_checkpoint(self): |
| 98 | # test single node |
| 99 | try: |
| 100 | tmpdir = tempfile.mkdtemp() |
| 101 | |
| 102 | def builder(): |
| 103 | ws = workspace.C.Workspace() |
| 104 | session = LocalSession(ws) |
| 105 | checkpoint = CheckpointManager(tmpdir, 'temp_node', 'minidb') |
| 106 | return session, checkpoint |
| 107 | |
| 108 | self.run_with(builder) |
| 109 | finally: |
| 110 | shutil.rmtree(tmpdir) |
| 111 | |
| 112 | # test multi-node |
| 113 | try: |
| 114 | tmpdir = tempfile.mkdtemp() |
| 115 | |
| 116 | def builder(): |
| 117 | ws = workspace.C.Workspace() |
| 118 | session = LocalSession(ws) |
| 119 | checkpoint = MultiNodeCheckpointManager(tmpdir, 'minidb') |
| 120 | return session, checkpoint |
| 121 | |
| 122 | self.run_with(builder) |
| 123 | finally: |
| 124 | shutil.rmtree(tmpdir) |
| 125 | |
| 126 | def test_ckpt_name_and_load_model_from_ckpts(self): |
| 127 | try: |