mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
import itertools
|
|
import numpy as np
|
|
import tempfile
|
|
import unittest
|
|
import os
|
|
|
|
from caffe2.python import core, workspace
|
|
import caffe2.python.hypothesis_test_util as hu
|
|
|
|
|
|
class TestMap(hu.HypothesisTestCase):
|
|
|
|
def test_create_map(self):
|
|
dtypes = [core.DataType.INT32, core.DataType.INT64]
|
|
for key_dtype, value_dtype in itertools.product(dtypes, dtypes):
|
|
op = core.CreateOperator(
|
|
'CreateMap',
|
|
[],
|
|
['map'],
|
|
key_dtype=key_dtype,
|
|
value_dtype=value_dtype,
|
|
)
|
|
workspace.RunOperatorOnce(op)
|
|
self.assertTrue(workspace.HasBlob('map'))
|
|
|
|
def test_map(self):
|
|
|
|
def test_map_func(KEY_T, VALUE_T):
|
|
model_file = os.path.join(tempfile.mkdtemp(), 'db')
|
|
key_data = np.asarray([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=KEY_T)
|
|
value_data = np.asarray([2, 3, 3, 3, 3, 2, 3, 3, 3, 3], dtype=VALUE_T)
|
|
workspace.FeedBlob("key_data", key_data)
|
|
workspace.FeedBlob("value_data", value_data)
|
|
save_net = core.Net("save_net")
|
|
save_net.KeyValueToMap(["key_data", "value_data"], "map_data")
|
|
save_net.Save(
|
|
["map_data"], [],
|
|
db=model_file,
|
|
db_type="minidb",
|
|
absolute_path=True
|
|
)
|
|
workspace.RunNetOnce(save_net)
|
|
workspace.ResetWorkspace()
|
|
load_net = core.Net("load_net")
|
|
load_net.Load(
|
|
[], ["map_data"],
|
|
db=model_file,
|
|
db_type="minidb",
|
|
load_all=True,
|
|
absolute_path=True
|
|
)
|
|
load_net.MapToKeyValue("map_data", ["key_data", "value_data"])
|
|
workspace.RunNetOnce(load_net)
|
|
key_data2 = workspace.FetchBlob("key_data")
|
|
value_data2 = workspace.FetchBlob("value_data")
|
|
assert(set(zip(key_data, value_data)) == set(zip(key_data2, value_data2)))
|
|
|
|
test_map_func(np.int64, np.int64)
|
|
test_map_func(np.int64, np.int32)
|
|
test_map_func(np.int32, np.int32)
|
|
test_map_func(np.int32, np.int64)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|