pytorch

1

2

3

4

5

6
from caffe2.python import core, workspace
7
from caffe2.python.test_util import TestCase
8
import tempfile
9

10

11
class TestCounterOps(TestCase):
12

13
    def test_counter_ops(self):
14
        workspace.RunOperatorOnce(core.CreateOperator(
15
            'CreateCounter', [], ['c'], init_count=1))
16

17
        workspace.RunOperatorOnce(core.CreateOperator(
18
            'CountDown', ['c'], ['t1']))  # 1 -> 0
19
        assert not workspace.FetchBlob('t1')
20

21
        workspace.RunOperatorOnce(core.CreateOperator(
22
            'CountDown', ['c'], ['t2']))  # 0 -> -1
23
        assert workspace.FetchBlob('t2')
24

25
        workspace.RunOperatorOnce(core.CreateOperator(
26
            'CountUp', ['c'], ['t21']))  # -1 -> 0
27
        assert workspace.FetchBlob('t21') == -1
28
        workspace.RunOperatorOnce(core.CreateOperator(
29
            'RetrieveCount', ['c'], ['t22']))
30
        assert workspace.FetchBlob('t22') == 0
31

32
        workspace.RunOperatorOnce(core.CreateOperator(
33
            'ResetCounter', ['c'], [], init_count=1))  # -> 1
34
        workspace.RunOperatorOnce(core.CreateOperator(
35
            'CountDown', ['c'], ['t3']))  # 1 -> 0
36
        assert not workspace.FetchBlob('t3')
37

38
        workspace.RunOperatorOnce(core.CreateOperator(
39
            'ResetCounter', ['c'], ['t31'], init_count=5))  # 0 -> 5
40
        assert workspace.FetchBlob('t31') == 0
41
        workspace.RunOperatorOnce(core.CreateOperator(
42
            'ResetCounter', ['c'], ['t32']))  # 5 -> 0
43
        assert workspace.FetchBlob('t32') == 5
44

45
        workspace.RunOperatorOnce(core.CreateOperator(
46
            'ConstantFill', [], ['t4'], value=False, shape=[],
47
            dtype=core.DataType.BOOL))
48
        assert workspace.FetchBlob('t4') == workspace.FetchBlob('t1')
49

50
        workspace.RunOperatorOnce(core.CreateOperator(
51
            'ConstantFill', [], ['t5'], value=True, shape=[],
52
            dtype=core.DataType.BOOL))
53
        assert workspace.FetchBlob('t5') == workspace.FetchBlob('t2')
54

55
        assert workspace.RunOperatorOnce(core.CreateOperator(
56
            'And', ['t1', 't2'], ['t6']))
57
        assert not workspace.FetchBlob('t6')  # True && False
58

59
        assert workspace.RunOperatorOnce(core.CreateOperator(
60
            'And', ['t2', 't5'], ['t7']))
61
        assert workspace.FetchBlob('t7')  # True && True
62

63
        workspace.RunOperatorOnce(core.CreateOperator(
64
            'CreateCounter', [], ['serialized_c'], init_count=22))
65
        with tempfile.NamedTemporaryFile() as tmp:
66
            workspace.RunOperatorOnce(core.CreateOperator(
67
                'Save', ['serialized_c'], [], absolute_path=1,
68
                db_type='minidb', db=tmp.name))
69
            for i in range(10):
70
                workspace.RunOperatorOnce(core.CreateOperator(
71
                    'CountDown', ['serialized_c'], ['t8']))
72
            workspace.RunOperatorOnce(core.CreateOperator(
73
                'RetrieveCount', ['serialized_c'], ['t8']))
74
            assert workspace.FetchBlob('t8') == 12
75
            workspace.RunOperatorOnce(core.CreateOperator(
76
                'Load', [], ['serialized_c'], absolute_path=1,
77
                db_type='minidb', db=tmp.name))
78
            workspace.RunOperatorOnce(core.CreateOperator(
79
                'RetrieveCount', ['serialized_c'], ['t8']))
80
            assert workspace.FetchBlob('t8') == 22
81

82
if __name__ == "__main__":
83
    import unittest
84
    unittest.main()
85