ideep_test_util.py 998 B

123456789101112131415161718192021222324252627282930313233343536373839
  1. ## @package ideep_test_util
  2. # Module caffe2.python.ideep_test_util
  3. """
  4. The IDEEP test utils is a small addition on top of the hypothesis test utils
  5. under caffe2/python, which allows one to more easily test IDEEP related
  6. operators.
  7. """
  8. import hypothesis.strategies as st
  9. from caffe2.proto import caffe2_pb2
  10. from caffe2.python import hypothesis_test_util as hu
  11. cpu_do = hu.cpu_do
  12. ideep_do = caffe2_pb2.DeviceOption(device_type=caffe2_pb2.IDEEP)
  13. device_options = hu.device_options + ([ideep_do])
  14. def device_checker_device_options():
  15. return st.just(device_options)
  16. def gradient_checker_device_option():
  17. return st.sampled_from(device_options)
  18. gcs = dict(
  19. gc=gradient_checker_device_option(),
  20. dc=device_checker_device_options()
  21. )
  22. gcs_cpu_only = dict(gc=st.sampled_from([cpu_do]), dc=st.just([cpu_do]))
  23. gcs_ideep_only = dict(gc=st.sampled_from([ideep_do]), dc=st.just([ideep_do]))
  24. gcs_cpu_ideep = dict(gc=st.sampled_from([cpu_do, ideep_do]), dc=st.just([cpu_do, ideep_do]))