def test_scheduler(self):
for scheduler in ('SIMPLE', 'FUSION', 'KNOWN', 'SIMPLE'):
try:
dragon.autograph.set_scheduler(scheduler)
if scheduler == 'FUSION':
self.assertEqual(config.config().graph_type, 'FusionGraph')
else:
self.assertEqual(config.config().graph_type, '')
except ValueError:
pass
def test_device(self):
major, minor = dragon.cuda.get_device_capability(0)
self.assertGreaterEqual(major, 1 if TEST_CUDA else 0)
self.assertGreaterEqual(minor, 0)
dragon.cuda.set_device(0)
self.assertEqual(dragon.cuda.current_device(), 0)
dragon.cuda.set_default_device(1)
self.assertEqual(config.config().device_type, 'cuda')
self.assertEqual(config.config().device_index, 1)
dragon.cuda.set_default_device(-1)
self.assertEqual(config.config().device_type, 'cpu')
self.assertEqual(config.config().device_index, 0)
def set_scheduler(scheduler='SIMPLE'):
"""Set the scheduler for symbolic graph.
Parameters
----------
scheduler : {'SIMPLE', 'FUSION'}, optional
The scheduler type.
"""
if scheduler not in ('SIMPLE', 'FUSION'):
raise ValueError('Unsupported scheduler: ' + scheduler)
if scheduler == 'SIMPLE':
config.config().graph_type = ''
elif scheduler == 'FUSION':
config.config().graph_type = 'FusionGraph'
def test_execution(self):
for mode in ('EAGER_MODE', 'GRAPH_MODE', 'UNKNOWN'):
try:
dragon.autograph.set_execution(mode)
self.assertEqual(config.config().graph_execution, mode)
except ValueError:
pass
def _add_device(graph_def):
"""Add device."""
cfg = config.config()
spec = context.get_device()
graph_def.device_option.CopyFrom(
proto_util.get_device_option(
spec.type, spec.index, cfg.random_seed))
def get_global_device_option():
"""Return the global device option."""
cfg = config.config()
return get_device_option(
cfg.device_type,
cfg.device_index,
)
def __init__(self, cache_key, device, **kwargs):
self._def = None
self._cache_key = cache_key
self._device = device
self._arg_device = proto_util.get_device_option('cpu')
self._arg_device = self._arg_device.SerializeToString()
self._seed = kwargs.get('seed', config.config().random_seed)
def run_backward(
self,
op_defs,
targets,
sources=None,
input_grads=None,
empty_grads=None,
):
"""Compute the gradients of input operators.
Parameters
----------
op_defs : Sequence[OperatorDef]
The executed op defs.
targets : Sequence[str]
The derivative targets.
sources : Sequence[str], optional
The differentiated inputs.
input_grads : Sequence[str], optional
The input grad for targets.
empty_grads : Sequence[str], optional
The grads to set to empty.
"""
cfg = config.config()
self.RunBackward(
op_defs,
targets,
sources if sources else [],
input_grads if input_grads else [],
empty_grads if empty_grads else [],
cfg.graph_optimization <= 2,
cfg.graph_verbosity > 0,
)
def get_device(use_default=True):
"""Return the nesting or default device."""
spec = _GLOBAL_DEVICE_STACK.get_default()
if spec is None:
cfg = config.config()
spec = device_spec.DeviceSpec(cfg.device_type, cfg.device_index)
return spec
def add_optimization(graph_def, level=None):
"""Add the optimization argument."""
cfg = config.config()
if level is None:
level = cfg.graph_optimization
graph_def.arg.add().CopyFrom(
proto_util.make_argument('optimization', level))
graph_def.graph_type = cfg.graph_type
def create_graph(self, graph_def):
"""Create a graph."""
cfg = config.config()
if cfg.graph_verbosity == 2:
msg = '\n' + str(graph_def)[:-1]
logging.info('\ngraph {' + msg.replace('\n', '\n ') + '\n}\n')
return self._impl.CreateGraph(serialization.serialize_proto(graph_def),
cfg.graph_verbosity == 1)
def reset_parameters(self):
"""Reset the parameters."""
numpy.random.seed(config.config().random_seed)
for li, di, pi in itertools.product(
range(self.num_layers),
range(self.num_directions),
range(self.num_gates * 2),
):
self.reset_parameter(li, di, pi, 'matrix', 'orthogonal')
self.reset_parameter(li, di, pi, 'bias', 'zero')
def set_directory(path):
"""Set the directory for logging files.
Parameters
----------
path : str, optional
The path of the directory.
"""
config.config().log_dir = path
def run_backward(self, op_defs, targets, grad_targets=None, sources=None):
"""Compute the gradients of operators."""
cfg = config.config()
self._impl.RunBackward(
op_defs,
targets,
grad_targets if grad_targets else [],
sources if sources else [],
cfg.graph_optimization > 2,
cfg.graph_verbosity > 0,
)
def add_device_option(graph_def):
"""Add the device option."""
cfg = config.config()
spec = context.get_device_spec()
graph_def.device_option.CopyFrom(
dragon_pb2.DeviceOption(device_type={
'cpu': 0,
'cuda': 1,
'cnml': 2
}[spec.type],
device_id=spec.index,
random_seed=cfg.random_seed))
def run_operator(self, op_def):
"""Run the operator.
Parameters
----------
op_def : Union[OperatorDef, Sequence[OperatorDef]]
The ``OperatorDef`` protocol buffer.
"""
cfg = config.config()
if isinstance(op_def, dragon_pb2.OperatorDef):
op_def = op_def.SerializePartialToString()
self.RunOperator(op_def, cfg.graph_verbosity > 0)
def set_default_device(device_index=0):
"""Set the default device.
A valid device index should be greater equal than 0:
```python
dragon.cuda.set_default_device(0) # Ok
dragon.cuda.set_default_device(-1) # Reset to the cpu device
```
Parameters
----------
device_index : int
The device index.
"""
if device_index < 0:
config.config().device_type = 'cpu'
config.config().device_index = 0
else:
config.config().device_type = 'cuda'
config.config().device_index = device_index
def __init__(self, **kwargs):
"""Create a ``DataTransformer``.
Parameters
----------
resize : int, optional, default=0
The size for the shortest edge.
padding : int, optional, default=0
The size for zero padding on two sides.
fill_value : int or sequence, optional, default=127
The value(s) to fill for padding or cutout.
crop_size : int, optional, default=0
The size for random-or-center cropping.
random_crop_size: int, optional, default=0
The size for sampling-based random cropping.
cutout_size : int, optional, default=0
The square size for the cutout algorithm.
mirror : bool, optional, default=False
Whether to apply the mirror (flip horizontally).
random_scales : Sequence[float], optional, default=(0.08, 1.)
The range of scales to sample a crop randomly.
random_aspect_ratios : Sequence[float], optional, default=(0.75, 1.33)
The range of aspect ratios to sample a crop randomly.
distort_color : bool, optional, default=False
Whether to apply color distortion.
inverse_color : bool, option, default=False
Whether to inverse channels for color images.
phase : {'TRAIN', 'TEST'}, optional
The optional running phase.
seed : int, optional
The random seed to use instead.
"""
super(DataTransformer, self).__init__()
self._resize = kwargs.get('resize', 0)
self._padding = kwargs.get('padding', 0)
self._fill_value = kwargs.get('fill_value', 127)
self._crop_size = kwargs.get('crop_size', 0)
self._random_crop_size = kwargs.get('random_crop_size', 0)
self._cutout_size = kwargs.get('cutout_size', 0)
self._mirror = kwargs.get('mirror', False)
self._random_scales = kwargs.get('random_scales', (0.08, 1.))
self._random_ratios = kwargs.get('random_aspect_ratios',
(3. / 4., 4. / 3.))
self._distort_color = kwargs.get('distort_color', False)
self._inverse_color = kwargs.get('inverse_color', False)
self._phase = kwargs.get('phase', 'TRAIN')
self._seed = kwargs.get('seed', config.config().random_seed)
self.q_in = self.q_out = None
self.daemon = True
def get_device_spec():
"""Return the device spec in current nesting."""
dev_info = get_device_info()
if dev_info is not None:
return device_spec.DeviceSpec(
dev_info['device_type'],
dev_info['device_index'],
)
else:
cfg = config.config()
return device_spec.DeviceSpec(
cfg.device_type,
cfg.device_index,
)
def set_verbosity(level=1):
"""Set the verbosity for graph ir.
Following levels are defined (default=0):
* level = ``0``: Do nothing.
* level = ``1``: Print the optimized GraphIR.
* level = ``2``: Print the raw GraphIR.
Parameters
----------
level : int, optional, default=1
The verbosity level.
"""
config.config().graph_verbosity = level
def __init__(self, key, dev, **kwargs):
"""Create a ``Function``.
Parameters
----------
key : str
The cache key.
device : dragon.vm.torch.device
The device spec.
"""
super(Function, self).__init__()
self._def = None
self._cache_key = key
self._device = dev
self._arg_device = proto_util.get_device_option('cpu')
self._arg_device = self._arg_device.SerializeToString()
self._seed = kwargs.get('seed', config.config().random_seed)
def set_optimization(level=1):
"""Set the optimization for graph ir.
Following levels are defined (default=3):
* level = ``0``: Do nothing.
* level = ``1``: Eliminate the unused outputs and operators.
* level = ``2``: Apply the inplace to inputs if available.
* level = ``3``: Allocate the shared buffer to outputs if available.
Parameters
----------
level : int, optional, default=3
The optimization level.
"""
config.config().graph_optimization = level
def cuda(self, device=None):
"""Switch the buffers and parameters to cuda device.
If ``device`` is not provided, use the value
set by ``dragon.config.set_cuda_device()``.
Parameters
----------
device : int, optional
The optional device index.
Returns
-------
dragon.vm.torch.nn.Module
The self.
"""
if device is None:
device = config.config().device_index
return self._apply(lambda t: t.cuda(device))
def create_graph(self, graph_def):
"""Create the graph.
Parameters
----------
graph_def : GraphDef
The ``GraphDef`` protocol buffer.
Returns
-------
str
The graph name.
"""
cfg = config.config()
if cfg.graph_verbosity == 2:
msg = '\n' + str(graph_def)[:-1]
logging.info('\ngraph {' + msg.replace('\n', '\n ') + '\n}\n')
return self.CreateGraph(
serialization.serialize_proto(graph_def),
cfg.graph_verbosity == 1)
def set_execution(execution='GRAPH_MODE'):
"""Set the execution mode for graph ir.
For changing the execution temporarily, use:
```python
# Enter a context to enforce graph execution
with dragon.graph_mode():
pass
# Enter a context to enforce eager execution
with dragon.eager_mode():
pass
```
Parameters
----------
execution : {'GRAPH_MODE', 'EAGER_MODE'}, optional
The execution mode.
"""
if execution not in ('GRAPH_MODE', 'EAGER_MODE'):
raise ValueError('Unsupported execution: ' + execution)
config.config().graph_execution = execution
def __init__(self, **kwargs):
"""Create a ``DataReader``.
Parameters
----------
dataset : class
The dataset class to load examples.
source : str
The path of data source.
part_idx : int, optional, default=0
The index of partition to read.
num_parts : int, optional, default=1
The total number of partitions over dataset.
shuffle : bool, optional, default=False
Whether to shuffle the data.
initial_fill : int, optional, default=1024
The length of sampling sequence for shuffle.
seed : int, optional
The random seed to use instead.
"""
super(DataReader, self).__init__()
self._dataset = kwargs.get('dataset', None)
self._source = kwargs.get('source', '')
self._part_idx = kwargs.get('part_idx', 0)
self._num_parts = kwargs.get('num_parts', 1)
self._shuffle = kwargs.get('shuffle', False)
self._initial_fill = kwargs.get('initial_fill', 1024) if self._shuffle else 1
self._seed = kwargs.get('seed', config.config().random_seed)
self._first, self._cursor, self._last = 0, 0, 0
self._part_size = 0
self._num_examples = 0
self._example_buffer = []
self._parts = []
self.q_out = None
self.daemon = True
def test_verbosity(self):
dragon.autograph.set_verbosity(1)
self.assertEqual(config.config().graph_verbosity, 1)
dragon.autograph.set_verbosity(0)
def run_operator(self, op_def):
"""Run an operator."""
cfg = config.config()
if isinstance(op_def, dragon_pb2.OperatorDef):
op_def = op_def.SerializePartialToString()
self._impl.RunOperator(op_def, cfg.graph_verbosity > 0)
def __init__(self):
super(_ThreadLocalData, self).__init__()
cfg = config.config()
self.mode = cfg.graph_execution
self.is_eager = self.mode == 'EAGER_MODE'
def test_optimization(self):
dragon.autograph.set_optimization(1)
self.assertEqual(config.config().graph_optimization, 1)
