transformations.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. # Copyright (c) 2016-present, Facebook, Inc.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. ##############################################################################
  15. import caffe2.python._import_c_extension as C
  16. class Transformer(object):
  17. def __init__(self):
  18. pass
  19. @classmethod
  20. def runTransform(cls, transform_name, net):
  21. pb = net.Proto().SerializeToString()
  22. if C.transform_exists(transform_name):
  23. output = C.run_transform(transform_name, pb)
  24. elif C.workspace_transform_exists(transform_name):
  25. output = C.run_workspace_transform(transform_name, pb)
  26. else:
  27. raise AttributeError('Transformation {} not found.'.format(transform_name))
  28. net.Proto().ParseFromString(output)
  29. def __getattr__(self, transform_name):
  30. return lambda net : self.runTransform(transform_name, net)
  31. def fuseNNPACKConvRelu(net):
  32. net.Proto().ParseFromString(
  33. C.transform_fuseNNPACKConvRelu(net.Proto().SerializeToString())
  34. )
  35. def optimizeForMKLDNN(net, training_mode = False):
  36. net.Proto().ParseFromString(
  37. C.transform_optimizeForMKLDNN(net.Proto().SerializeToString(), training_mode)
  38. )
  39. def fuseConvBN(net):
  40. net.Proto().ParseFromString(
  41. C.transform_fuseConvBN(net.Proto().SerializeToString())
  42. )