pytorch

1

2

3

4

5

6
from caffe2.python import core, workspace
7
from caffe2.python.test_util import TestCase
8
import numpy as np
9

10

11
class TestCounterOps(TestCase):
12

13
    def test_stats_ops(self):
14
        # The global StatRegistry isn't reset when the workspace is reset,
15
        #   so there may be existing stats from a previous test
16
        workspace.RunOperatorOnce(core.CreateOperator(
17
            'StatRegistryExport', [], ['prev_k', 'prev_v', 'prev_ts']))
18
        previous_keys = workspace.FetchBlob('prev_k')
19
        existing = len(previous_keys)
20

21
        prefix = '/'.join([__name__, 'TestCounterOps', 'test_stats_ops'])
22
        keys = [
23
            (prefix + '/key1').encode('ascii'),
24
            (prefix + '/key2').encode('ascii')
25
        ]
26
        values = [34, 45]
27
        workspace.FeedBlob('k', np.array(keys, dtype=str))
28
        workspace.FeedBlob('v', np.array(values, dtype=np.int64))
29
        for _ in range(2):
30
            workspace.RunOperatorOnce(core.CreateOperator(
31
                'StatRegistryUpdate', ['k', 'v'], []))
32
        workspace.RunOperatorOnce(core.CreateOperator(
33
            'StatRegistryExport', [], ['k2', 'v2', 't2']))
34

35
        workspace.RunOperatorOnce(core.CreateOperator(
36
            'StatRegistryCreate', [], ['reg']))
37
        workspace.RunOperatorOnce(core.CreateOperator(
38
            'StatRegistryUpdate', ['k2', 'v2', 'reg'], []))
39

40
        workspace.RunOperatorOnce(core.CreateOperator(
41
            'StatRegistryExport', ['reg'], ['k3', 'v3', 't3']))
42

43
        k3 = workspace.FetchBlob('k3')
44
        v3 = workspace.FetchBlob('v3')
45
        t3 = workspace.FetchBlob('t3')
46

47
        self.assertEqual(len(k3) - existing, 2)
48
        self.assertEqual(len(v3), len(k3))
49
        self.assertEqual(len(t3), len(k3))
50
        for key in keys:
51
            self.assertIn(key, k3)
52