pytorch

1

2

3

4

5

6
from caffe2.python import core, workspace, test_util
7
import os
8
import shutil
9
import tempfile
10
import unittest
11

12

13
class CheckpointTest(test_util.TestCase):
14
    """A simple test case to make sure that the checkpoint behavior is correct.
15
    """
16

17
    @unittest.skipIf("LevelDB" not in core.C.registered_dbs(), "Need LevelDB")
18
    def testCheckpoint(self):
19
        temp_root = tempfile.mkdtemp()
20
        net = core.Net("test_checkpoint")
21
        # Note(jiayq): I am being a bit lazy here and am using the old iter
22
        # convention that does not have an input. Optionally change it to the
23
        # new style if needed.
24
        net.Iter([], "iter")
25
        net.ConstantFill([], "value", shape=[1, 2, 3])
26
        net.Checkpoint(["iter", "value"], [],
27
                       db=os.path.join(temp_root, "test_checkpoint_at_%05d"),
28
                       db_type="leveldb", every=10, absolute_path=True)
29
        self.assertTrue(workspace.CreateNet(net))
30
        for i in range(100):
31
            self.assertTrue(workspace.RunNet("test_checkpoint"))
32
        for i in range(1, 10):
33
            # Print statements are only for debugging purposes.
34
            # print("Asserting %d" % i)
35
            # print(os.path.join(temp_root, "test_checkpoint_at_%05d" % (i * 10)))
36
            self.assertTrue(os.path.exists(
37
                os.path.join(temp_root, "test_checkpoint_at_%05d" % (i * 10))))
38

39
        # Finally, clean up.
40
        shutil.rmtree(temp_root)
41

42

43
if __name__ == "__main__":
44
    unittest.main()
45