utils_test.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. from caffe2.python import core, utils, test_util
  2. import numpy as np
  3. class TestUtils(test_util.TestCase):
  4. def testArgsToDict(self):
  5. args = [utils.MakeArgument("int1", 3),
  6. utils.MakeArgument("float1", 4.0),
  7. utils.MakeArgument("string1", "foo"),
  8. utils.MakeArgument("intlist1", np.array([3, 4])),
  9. utils.MakeArgument("floatlist1", np.array([5.0, 6.0])),
  10. utils.MakeArgument("stringlist1", np.array(["foo", "bar"]))]
  11. dict_ = utils.ArgsToDict(args)
  12. expected = {"int1" : 3,
  13. "float1" : 4.0,
  14. "string1" : b"foo",
  15. "intlist1" : [3, 4],
  16. "floatlist1" : [5.0, 6.0],
  17. "stringlist1" : [b"foo", b"bar"]}
  18. self.assertEqual(dict_, expected, "dictionary version of arguments "
  19. "doesn't match original")
  20. def testBuildUniqueMutexIter(self):
  21. init_net = core.Net("init_net")
  22. net = core.Net("net")
  23. utils.BuildUniqueMutexIter(init_net, net)
  24. for op in init_net.Proto().op:
  25. self.assertEqual(op.device_option.extra_info[0],
  26. "device_type_override:cpu")
  27. for op in net.Proto().op:
  28. self.assertEqual(op.device_option.extra_info[0],
  29. "device_type_override:cpu")