1234567891011121314151617181920212223242526 |
- """
- Copyright 2020 The Microsoft DeepSpeed Team
- """
- import os
- import pkgutil
- import importlib
- from deepspeed.accelerator import get_accelerator
- # List of all available ops
- # reflect all builder names into __op_builders__
- op_builder_dir = get_accelerator().op_builder_dir()
- op_builder_module = importlib.import_module(op_builder_dir)
- __op_builders__ = []
- for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]):
- # avoid self references
- if module_name != 'all_ops' and module_name != 'builder' and module_name != 'builder_names':
- module = importlib.import_module("{}.{}".format(op_builder_dir, module_name))
- for member_name in module.__dir__():
- if member_name.endswith('Builder'):
- # append builder to __op_builders__ list
- builder = get_accelerator().create_op_builder(member_name)
- __op_builders__.append(builder)
- ALL_OPS = {op.name: op for op in __op_builders__ if op is not None}
|