def normalize_spatial_args(name, values, num_spatial_dims):
if name in ('ksize', 'strides', 'dilations'):
if values is None:
return [1] * num_spatial_dims
else:
values = nest.flatten(values)
if len(values) == 1:
return [values[0]] * num_spatial_dims
elif len(values) != num_spatial_dims:
defaults = [1] * num_spatial_dims
defaults[:num_spatial_dims] = values
return defaults
return values
elif name == 'padding':
if isinstance(values, six.string_types):
padding, pads = values.upper(), 0
else:
padding_tuple = nest.flatten(values)
padding = 'VALID'
if len(padding_tuple) == 1:
pads = padding_tuple[0]
elif len(padding_tuple) == num_spatial_dims:
pads = padding_tuple
elif len(padding_tuple) == (num_spatial_dims * 2):
pads_l, pads_r = [], []
for i in range(num_spatial_dims):
pads_l.append(padding_tuple[i * 2])
pads_r.append(padding_tuple[i * 2 + 1])
pads = pads_l + pads_r
else:
raise ValueError(
'Except 1, {} or {} values if <padding> set as explict pads.'
.format(num_spatial_dims, num_spatial_dims * 2))
return padding, pads
def _add_node(self, inputs, outputs):
"""Add a layer node for inputs and outputs.
Parameters
----------
inputs : Sequence[dragon.Tensor]
The input tensors.
outputs : Sequence[dragon.Tensor]
The output tensors.
"""
inputs = nest.flatten(inputs)
outputs = nest.flatten(outputs)
input_info = [getattr(e, '_info', [None, None]) for e in inputs]
self._nodes.append(
node.LayerNode(
self,
node_index=len(self._nodes),
in_nodes=[e[0] for e in input_info],
in_tensor_idxes=[e[1] for e in input_info],
in_tensors=inputs,
out_tensors=outputs,
))
for idx, tensor in enumerate(outputs):
tensor._info = (self._nodes[-1], idx)
def assertEqual(
self,
first,
second,
msg=None,
prec=None,
):
if prec is None:
prec = self.precision
inputs = nest.flatten(first)
num_first = len(inputs)
inputs += nest.flatten(second)
num_second = len(inputs) - num_first
for i, input in enumerate(inputs):
if isinstance(input, torch.Tensor):
inputs[i] = input.numpy()
first = inputs[:num_first] if num_first > 1 else inputs[0]
second = inputs[num_first:len(inputs)] if num_second > 1 else inputs[
num_first]
if isinstance(first, np.ndarray) and isinstance(second, np.ndarray):
super(OpTestCase, self).assertEqual(first.shape, second.shape)
if first.dtype == np.bool and second.dtype == np.bool:
diff = first ^ second
num_unique = len(np.unique(diff))
self.assertLessEqual(num_unique, 1, msg)
else:
diff = np.abs(first - second)
max_err = diff.max()
self.assertLessEqual(max_err, prec, msg)
elif nest.is_sequence(first) and nest.is_sequence(second):
for a, b in zip(first, second):
self.assertEqual(a, b, msg, prec)
else:
super(OpTestCase, self).assertEqual(first, second, msg)
def linspace(start, stop, num, dtype='int64', axis=0, **kwargs):
r"""Generate evenly spaced values within intervals along the given axis.
Range :math:`[\text{start}, \text{stop})` is determined for :attr:`num` values:
```python
x = dragon.linspace(2, 4, num=3) # [2, 3, 4]
```
More than one ranges are accepted to generate N-d coordinates:
```python
x = dragon.linspace([1, 2], [3, 4], num=3, axis=0) # [[1, 2], [2, 3], [3, 4]]
y = dragon.linspace([1, 2], [3, 4], num=3, axis=1) # [[1, 2, 3], [2, 3, 4]]
```
Parameters
----------
start : Union[number, Sequence[number]]
The start of range.
stop: Union[number, Sequence[number]]
The stop of range.
num : int
The number of values to generate.
dtype : str, optional, default='int64'
The optional data type.
axis : int, optional, default=0
The axis to generate values.
Returns
-------
dragon.Tensor
The output tensor.
"""
dtype = dtype.lower()
starts = [float(elem) for elem in nest.flatten(start)]
stops = [float(elem) for elem in nest.flatten(stop)]
dims = []
if len(starts) > 1 or starts == start:
dims = [len(starts)]
axis = axis if axis >= 0 else axis + len(dims) + 1
dims.insert(axis, num)
if context.executing_eagerly():
return OpLib.execute('LinSpace', [],
ndim=len(dims),
num_intervals=len(starts),
axis=axis,
dtype=dtype,
dims=dims,
start=starts,
stop=stops)
return OpLib.add('LinSpace', [],
axis=axis,
dtype=dtype,
dims=dims,
start=starts,
stop=stops,
**kwargs)
def add(op_type, inputs, **kwargs):
"""Add operator to output symbols."""
op_tape = tapes.OrderedTape()
graph_tape = tapes.get_tape()
execute_ws = workspace.get_workspace()
# Add inputs.
enable_grad = False
inputs = nest.flatten(inputs)
for input in inputs:
op_tape.add_source(input)
if graph_tape and (input.requires_grad
or graph_tape.is_target(id(input))):
enable_grad = True
# Add extra inputs.
for input in nest.flatten(kwargs.pop('extra_inputs', [])):
op_tape.add_source(input)
op_tape.add_target(input.id)
# Add outputs.
name = kwargs.pop('name', None)
num_outputs = kwargs.pop('num_outputs', 1)
outputs = []
for i in range(num_outputs):
outputs.append(
Tensor(impl=execute_ws.create_tensor(scope='Tensor'),
name=name if name else op_type + ':%d' % i,
symbolic=True))
# Create def.
op_def = proto_util.make_operator_def(
op_type=op_type,
inputs=[input.id for input in inputs],
outputs=[output.id for output in outputs],
device_option=proto_util.get_default_device_option(),
name=execute_ws.create_handle('Op'),
**kwargs)
# Record def.
op_tape.add_element(op_def)
graph_tape.add_element(op_def) if enable_grad else None
# Set tape for outputs.
for output in outputs:
output._tape = op_tape
output._requires_grad = enable_grad
# Add spec to outputs.
add_output_spec = OpSchema.get_spec(op_type)
if add_output_spec is None:
add_output_spec = OpSchema.get_spec('Unchanged')
outputs = add_output_spec(kwargs, inputs, outputs)
# Return single or repeated outputs.
return outputs[0] if num_outputs == 1 else outputs
def _init_graph_network(self, inputs, outputs, **kwargs):
self._is_graph_network = True
if isinstance(inputs, list) and len(nest.flatten(inputs)) == 1:
inputs = inputs[0]
if isinstance(outputs, list) and len(nest.flatten(outputs)) == 1:
outputs = outputs[0]
self._nested_outputs = outputs
self._nested_inputs = inputs
self.inputs = nest.flatten(inputs)
self.outputs = nest.flatten(outputs)
self.built = True
def gradient(self, target, sources, output_gradients=None):
"""Compute the derivatives of target w.r.t. sources."""
# Check the pushed tape.
if self._tape is None:
raise RuntimeError('GradientTape.gradient(...) can only be called '
'once on non-persistent tapes.')
# Stop recording if not persistent.
if self._recording and not self._persistent:
self._pop_tape()
# Check the gradient of outputs.
xs, ys = nest.flatten(sources), nest.flatten(target)
grad_ys = []
if output_gradients is not None:
for y, grad_y in zip(ys, nest.flatten(output_gradients)):
if grad_y.shape != y.shape:
raise ValueError(
'Excepted the dimensions of output gradient is {}, '
'got {}.'.format(y.shape, grad_y.shape))
grad_ys.append(grad_y)
# Record or compute the gradient of inputs.
execute_ws = workspace.get_workspace()
grad_xs = []
if not context.executing_eagerly():
for x in xs:
self._tape.add_source(x)
grad_xs.append(
Tensor(shape=x.shape,
dtype=x.dtype,
impl=execute_ws.create_tensor(x.id + '_grad'),
symbolic=True))
for i, y in enumerate(ys):
y._grad_tape = self._tape
y._grad = grad_ys[i] if i < len(grad_ys) else None
else:
execute_ws.run_backward(op_defs=self._tape.get_elements(),
targets=[y.id for y in ys],
sources=[x.id for x in xs],
grad_targets=[dy.id for dy in grad_ys])
for x in xs:
impl = execute_ws.get_tensor(x.id + '_grad')
grad_xs.append(Tensor(impl=impl) if impl else None)
# Remove tape if not persistent.
if not self._persistent:
self._release_tape(execute_ws)
return grad_xs
def all_reduce(tensor, op='SUM', group=None):
"""Reduce the tensor across all nodes in a group.
Parameters
----------
tensor : Sequence[dragon.vm.torch.Tensor]
The tensor(s) to reduce.
op : {'SUM', 'MEAN'}, optional
The reduce operation.
group : ProcessGroup, optional
The group for communication.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
if group is None:
group = distributed.get_group()
if group is None:
raise ValueError('<group> is required.')
if op not in ('MEAN', 'SUM'):
raise ValueError('Unsupported reduce op: ' + op)
tensors = nest.flatten(tensor)
return _functions.Collective \
.instantiate(
tensors[0].device,
operation=op,
communication='ALLREDUCE',
group=group,
).apply(tensors)
def zeros(*size, out=None, dtype='float32', device=None, requires_grad=False):
r"""Return a tensor filled with zeros.
.. math:: \text{out} \leftarrow 0
Parameters
----------
size : int...
The output tensor shape.
out : dragon.vm.torch.Tensor, optional
The output tensor.
dtype : str, optional, default='float32'
The data type of output tensor.
device : dragon.vm.torch.device, optional
The device of output tensor.
requires_grad : bool, optional, default=False
Record gradient for output tensor or not.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
size = nest.flatten(size)
device = out.device if out else (device or cpp.device())
out = Function.apply('Fill',
device, [],
outputs=[out],
dtype=dtype,
value=0.0,
ndim=len(size),
dims=size)
out._requires_grad = requires_grad
return out
def _build_graphs(self, *args, **kwargs):
attributes = self._attribute_cache[workspace.get_workspace()]
input_signature = self._spec.input_signature
args, kwargs = self._spec.separate_inputs(*args, **kwargs)
inputs = []
for i in range(self._spec.num_inputs):
input_spec = None
if input_signature is not None:
input_spec = input_signature[i]
if not isinstance(args[i], Tensor) and input_spec is None:
inputs.append(args[i])
continue
name = 'Input_%d' % (i + 1)
shape = getattr(args[i], 'shape', None)
dtype = getattr(args[i], 'dtype', None)
if input_spec is not None:
shape, dtype = input_spec.shape, input_spec.dtype
inputs.append(Tensor(shape, dtype, name=name, symbolic=True))
with eager_context.graph_mode():
outputs = self._run_function(*inputs, **kwargs)
graph_outputs, dummies, graphs = [], [], []
for output in nest.flatten(outputs):
if isinstance(output, Tensor):
graph_outputs.append(output)
else:
dummies.append(output)
if len(graph_outputs) > 0:
graphs.append(GraphLib.from_outputs(graph_outputs))
for obj in dummies:
if isinstance(obj, GraphExec):
graphs.append(obj)
attributes['inputs'] = inputs
attributes['outputs'] = outputs
attributes['graphs'] = graphs
return graphs
def new_group(ranks=None, backend=None, verbose=False):
"""Create a new communication group.
The ``ranks`` can be set to **None** to create
an empty group for disabling the distributed utilities.
If ``backend`` is **None**, select as: **NCCL** > **MPI**.
Note that this function should be called from all processes,
even if they are not going to be included in this group.
Parameters
----------
ranks : Sequence[int], optional
The rank of processes to be included.
backend : {'AUTO', 'MPI', 'NCCL'}, optional
The optional backend.
verbose : bool, optional, default=False
**True** to log the group info.
"""
if ranks is None:
return ProcessGroup(None, None, None, backend)
else:
_maybe_initialize()
ranks = nest.flatten(ranks)
comm, handle = _b.mpiCreateGroup(ranks, verbose)
return ProcessGroup(ranks, comm, handle, backend)
def _setup(self):
"""Connect the layers sequentially."""
self._net_outputs = set()
# Collect bottom and top blobs.
for layer_idx, layer in enumerate(self._layers):
bottom = []
for blob in layer._bottom:
if blob not in self._blobs:
raise RuntimeError('bottom({}) is unknown.'.format(blob))
bottom.append(self._blobs[blob])
if blob in self._net_outputs:
self._net_outputs.remove(blob)
if isinstance(layer, layer_factory.BatchNorm):
next_layer = self._layers[layer_idx + 1]
if isinstance(next_layer, layer_factory.Scale):
layer.fuse_with_scale_layer(next_layer)
with context.name_scope(layer._name):
outputs = layer.setup([blob['data'] for blob in bottom])
if outputs is not None:
outputs = nest.flatten(outputs)
for blob_idx, blob in enumerate(layer._top):
self._blobs[blob] = {
'data': outputs[blob_idx],
'diff': TensorRef(outputs[blob_idx].id + '_grad')
}
self._net_outputs.add(blob)
# Collect layer param blobs.
for blobs in self.params.values():
self._layer_blobs.extend(blobs)
def flip(input, dims):
"""Reverse elements along the given dimension.
:attr:`dims` could be negative:
```python
x = torch.tensor([[1, 2, 3], [4, 5, 6]])
# A negative dimension is the last-k dimension
print(torch.flip(x, dims=1)) # [[3, 2, 1], [6, 5, 4]]
print(torch.flip(x, dims=-1)) # Equivalent
# Also, dimension could be a sequence of integers
print(torch.flip(x, dims=(0, 1))) # [[6, 5, 4], [3, 2, 1]]
```
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
dims : Union[int, Sequence[int]]
The dimension to reverse.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
return Function.apply(
'Reverse',
input.device, [input],
axes=nest.flatten(dims) if dims is not None else dims)
def index_select(input, dim, index, out=None):
"""Select the elements along the given dim using index.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
dim : Union[int, Sequence[int]]
The dim(s) to select.
index : dragon.vm.torch.Tensor
The index tensor.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dim = nest.flatten(dim)
dim.sort()
if dim[-1] != (dim[0] + len(dim) - 1):
raise ValueError('The <dim> should be a continuous sequence.')
return _functions.IndexSelect \
.instantiate(
utils.unify_devices([input, index]),
axis=dim[0],
num_axes=len(dim),
).apply(input, index, out)
def reverse(inputs, axis, **kwargs):
"""Reverse elements along the given axis.
:attr:`axis` could be negative:
```python
x = dragon.constant([[1, 2, 3], [4, 5, 6]])
# A negative axis is the last-k axis
print(dragon.reverse(x, axis=1)) # [[3, 2, 1], [6, 5, 4]]
print(dragon.reverse(x, axis=-1)) # Equivalent
# Also, axis could be a sequence of integers
print(dragon.reverse(x, axis=(0, 1))) # [[6, 5, 4], [3, 2, 1]]
```
Parameters
----------
inputs : dragon.Tensor
The input tensor.
axis : Union[int, Sequence[int]]
The axis to reverse.
Returns
-------
dragon.Tensor
The output tensor.
"""
axes = nest.flatten(axis) if axis is not None else axis
if context.executing_eagerly():
return OpLib.execute('Reverse', inputs, axes=axes)
return OpLib.add('Reverse', inputs, axes=axes, **kwargs)
def randn(*size, out=None, dtype='float32', device=None, requires_grad=False):
"""Return a tensor from the normal distribution of N(0, 1).
Parameters
----------
size : int...
The output tensor shape.
out : dragon.vm.torch.Tensor, optional
The output tensor.
dtype : str, optional, default='float32'
The data type of output tensor.
device : dragon.vm.torch.device, optional
The device of output tensor.
requires_grad : bool, optional, default=False
Record gradient for output tensor or not.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
size = nest.flatten(size)
device = out.device if out else (device or cpp.device())
out = Function.apply(
'RandomNormal', device, [], outputs=[out],
dtype=dtype, mean=0.0, std=1.0, ndim=len(size), dims=size)
out._requires_grad = requires_grad
return out
def forward(self, inputs, **kwargs):
"""Compute the output of RNN."""
inputs = nest.flatten(inputs)
op_lib = rnn_ops_lib.Recurrent
if context.executing_eagerly():
inputs.insert(1, self._weights)
return op_lib \
.instantiate(
mode=self.mode,
num_layers=self.num_layers,
hidden_size=self.hidden_size,
bidirectional=self.bidirectional,
dropout_ratio=self.dropout,
is_training=kwargs.get('is_training', False),
).apply(inputs)
else:
inputs.insert(1, self._weights_ref)
return op_lib.blend(
inputs=inputs,
rnn_mode=self.mode,
num_layers=self.num_layers,
hidden_size=self.hidden_size,
bidirectional=self.bidirectional,
dropout_ratio=self.dropout,
)
def broadcast(tensor, src=0, group=None):
"""Broadcast the tensor from source node in a group.
Parameters
----------
tensor : Sequence[dragon.vm.torch.Tensor]
The tensor(s) to reduce.
src : int
The rank of the source node.
group : ProcessGroup, optional
The group for communication.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
if group is None:
group = distributed.get_group()
if group is None:
raise ValueError('<group> is required.')
tensors = nest.flatten(tensor)
return _functions.Collective \
.instantiate(
tensors[0].device,
root=src,
communication='BROADCAST',
group=group,
).apply(tensors)
def normalize_tuple(value, rank):
"""Repeat the value according to the rank."""
value = nest.flatten(value)
if len(value) > rank:
return (value[i] for i in range(rank))
else:
return tuple([value[i] for i in range(len(value))] +
[value[-1] for _ in range(len(value), rank)])
def test_flatten(self):
self.assertEqual(nest.flatten(1), [1])
for a, b in zip(nest.flatten_with_paths([2, 4, [31, 32], 1]),
[((0,), 2), ((1,), 4), ((2, 0), 31), ((2, 1), 32), ((3,), 1)]):
self.assertEqual(a, b)
for a, b in zip(nest.flatten_with_paths({2: 2, 4: 4, 3: {1: 31, 2: 32}, 1: 1}),
[((1,), 1), ((2,), 2), ((3, 1), 31), ((3, 2), 32), ((4,), 4)]):
self.assertEqual(a, b)
def verify(inputs, min_num, max_num):
if min_num == max_num and min_num == 0:
return
inputs = nest.flatten(inputs)
if len(inputs) < min_num or len(inputs) > max_num:
raise ValueError('The number of <inputs> is {}, '
'not in range: [min={}, max={}].'.format(
len(inputs), min_num, max_num))
def gradient(self, target, sources, output_gradients=None):
"""Compute the derivatives of ``target`` w.r.t. ``sources``."""
# Fallback to the symbolic implementation.
if not context.executing_eagerly():
return grad_impl.gradients(
ys=target,
xs=sources,
grad_ys=output_gradients,
)
# Check the pushed tape.
if self._tape is None:
raise RuntimeError('GradientTape.gradient(...) can only be called '
'once on non-persistent tapes.')
# Stop recording if not persistent.
if self._recording:
if not self._persistent:
self._pop_tape()
# Collect gradient info.
xs, ys, grad_ys = nest.flatten(sources), nest.flatten(target), []
if output_gradients is not None:
for tensor, grad_tensor in zip(ys, nest.flatten(output_gradients)):
if grad_tensor.shape != tensor.shape:
raise ValueError(
'Excepted the dimensions of output gradient is {}, '
'got {}.'.format(tensor.shape, grad_tensor.shape))
grad_ys.append(grad_tensor.id)
# Run the gradient ops sequentially.
current_ws = workspace.get_workspace()
current_ws.run_backward(
op_defs=self._tape._defs,
targets=[y.id for y in ys],
sources=[x.id for x in xs],
input_grads=grad_ys,
empty_grads=self._tape.empty_grads,
)
# Remove the tape to release resources.
if not self._persistent:
self._tape = None
# Pack the gradients.
return [_steal_grad(current_ws, x) for x in xs]
def _normalize_tuple(value, rank):
"""Repeat the value to a tuple."""
value = nest.flatten(value)
if len(value) > rank:
return [value[i] for i in range(rank)]
else:
return [value[i] for i in range(len(value))] + \
[value[-1] for _ in range(len(value), rank)]
def roll(inputs, shift, axis=None, **kwargs):
"""Roll elements along the given axis.
:attr:`axis` could be negative or ``None``:
```python
x = dragon.constant([[1, 2, 3], [4, 5, 6]])
# A negative axis is the last-k axis
print(dragon.roll(x, shift=1, axis=1)) # [[3, 1, 2], [6, 4, 5]]
print(dragon.roll(x, shift=1, axis=-1)) # Equivalent
# If axis is None, roll input as a vector
print(dragon.roll(x, shift=1)) # [[6, 1, 2], [3, 4, 5]]
# Also, axis could be a sequence of integers
print(dragon.roll(x, shift=(1, 1), axis=(0, 1))) # [[6, 4, 5], [3, 1, 2]]
print(dragon.roll(x, shift=(1, -1), axis=(0, 1))) # [[5, 6, 4], [2, 3, 1]]
```
Parameters
----------
inputs : dragon.Tensor
The input tensor.
shift : Union[int, Sequence[int], dragon.Tensor]
The rolling offset of each axis.
axis : Union[int, Sequence[int]], optional
The axis to roll.
Returns
-------
dragon.Tensor
The output tensor.
"""
args = OpSchema.parse_args(locals())
axes = nest.flatten(axis) if axis is not None else axis
if isinstance(shift, six.integer_types):
args['shifts'] = nest.flatten(shift)
if context.executing_eagerly():
return OpLib.execute(
'Roll', inputs, num_shifts=len(args['shifts']),
shifts=args['shifts'], axes=axes)
args.pop('axis')
return OpLib.add('Roll', axes=axes, **args)
def _init(self):
"""Connect the layers sequentially."""
losses, learnable_blobs = [], []
grad_tape = backprop.GradientTape()
# Collect bottom and top blobs.
for i, layer in enumerate(self._layers):
bottoms = []
for bottom_name in layer.bottom:
if bottom_name not in self._net_blobs:
raise RuntimeError(
'Bottom "{}" is unknown.'.format(bottom_name))
bottoms.append(self._net_blobs[bottom_name])
if bottom_name in self._net_outputs:
self._net_outputs.remove(bottom_name)
if isinstance(layer, layer_factory.BatchNorm):
next_layer = self._layers[i + 1]
if isinstance(next_layer, layer_factory.Scale):
layer.scale_layer = next_layer
with context.name_scope(layer.name), grad_tape:
outputs = layer.setup([blob['data'] for blob in bottoms])
if outputs is not None:
outputs = nest.flatten(outputs)
for j, top_name in enumerate(layer.top):
self._net_blobs[top_name] = {
'data': outputs[j],
'diff': None
}
self._net_outputs.add(top_name)
loss_weights = list(layer._proto.loss_weight)
if layer._proto.type.find('Loss') != -1:
if len(loss_weights) == 0:
loss_weights.append(1)
for j, loss_weight in enumerate(loss_weights):
if loss_weight > 0:
losses.append(outputs[j])
for j, blob in enumerate(layer.blobs):
lr_mult, decay_mult = 1, 1
if j < len(layer._proto.param):
p = layer._proto.param[j]
lr_mult = p.lr_mult if p.HasField('lr_mult') else 1
decay_mult = p.decay_mult if p.HasField(
'decay_mult') else 1
if lr_mult > 0 and blob['data'].requires_grad:
if decay_mult == 0:
blob['data']._weight_decay = 0
learnable_blobs.append(blob)
if self._phase == 'TRAIN':
with eager_context.graph_mode():
grads = grad_tape.gradient(
losses, [blob['data'] for blob in learnable_blobs])
for blob, grad in zip(learnable_blobs, grads):
blob['diff'] = grad
# Collect all learnable blobs.
for blobs in self.params.values():
for blob in blobs:
if blob.diff:
self._learnable_blobs.append(blob)
def forward(self, input, weights, hx=None):
inputs = [input, weights]
if hx is not None:
inputs += nest.flatten(hx)
outputs = [self.alloc() for _ in range(self.num_outputs)]
outputs = self.dispatch(inputs, outputs)
if self.num_outputs == 3:
return outputs[0], outputs[1:]
return outputs
def moments(inputs, axis=None, keepdims=False, **kwargs):
r"""Compute the mean and variance of input along the given axis.
.. math::
\begin{cases}
\mathrm{E}[x] = \frac{1}{n}\sum(x) \\
\mathrm{Var}[x] = \frac{1}{n}\sum(x - \mathrm{E}[x])^{2}
\end{cases}
:attr:`axis` could be negative or ``None``:
```python
x = dragon.constant([[1, 2, 3], [4, 5, 6]], dtype='float32')
# A negative axis is the last-k axis
print(dragon.nn.moments(x, 1))
print(dragon.nn.moments(x, -1)) # Equivalent
# If axis is None, reduce as a vector and return scalars
# will be applied to return a scalar result
print(dragon.nn.moments(x)) # mean is 3.5, var is 2.916667
# Also, axis could be a sequence of integers
print(dragon.nn.moments(x, (0, 1))) # mean is 3.5, var is 2.916667
```
Parameters
----------
inputs : dragon.Tensor
The input tensor.
axis : Union[int, Sequence[int]], optional
The axis to reduce.
keepdims : bool, optional, default=False
Keep the reduced dimensions or not.
Returns
-------
dragon.Tensor
The mean tensor.
dragon.Tensor
The variance tensor.
"""
axes = None if axis is None else nest.flatten(axis)
if context.executing_eagerly():
return OpLib.execute('Moments',
inputs,
outputs=[None, None],
axes=axes,
keepdims=keepdims)
return OpLib.add('Moments',
inputs,
num_outputs=2,
axes=axes,
keepdims=keepdims,
**kwargs)
def __call__(self, *args, **kwargs):
"""Call the compiled executables."""
if self.executables is None:
# Graph is not created on the first call.
# Compile the executables from the python function.
inputs = []
input_signature = self.input_signature
with context.name_scope('${%d}' % id(self)):
for i in range(self._function_spec.num_inputs):
name, shape, dtype = 'Input:%d' % i, None, None
if input_signature is not None:
if i >= len(input_signature):
raise ValueError(
'When <input_signature> is provided, '
'only define arguments covered by it.\n'
'Got %d signature(s) and %d argument(s).' %
(len(input_signature),
self._function_spec.num_inputs))
shape = input_signature[i].shape
dtype = input_signature[i].dtype
inputs.append(Tensor(shape, dtype, name).constant())
with context.name_scope('${%d}' %
id(self)), eager_context.graph_mode():
returns = nest.flatten(self._python_function(*inputs))
outputs, dummies = [], []
for obj in returns:
if isinstance(obj, Tensor):
outputs.append(obj)
else:
dummies.append(obj)
executables = [function_lib.create_function(outputs=outputs)]
for obj in dummies:
if isinstance(obj, optimizer.Optimizer):
executables.append(
function_lib.create_function(optimizer=obj))
self.inputs = inputs
self.outputs = returns
self.executables = executables
# In this case, we have compiled executables.
# Notify the backend to run directly.
executables = self.executables
inputs, kwargs = self.canonicalize_inputs(*args, **kwargs)
current_ws = workspace.get_workspace()
for input, value in zip(self.inputs, inputs):
current_ws.feed_tensor(input, value)
executables[0](return_outputs=False, **kwargs)
[func(return_outputs=False) for func in executables[1:]]
outputs = []
for output in self.outputs:
if isinstance(output, Tensor):
impl = current_ws.GetTensor(output.id)
device = device_spec.DeviceSpec(*impl.device)
outputs.append(EagerTensor(impl=impl, device=device))
else:
outputs.append(output)
return outputs[0] if len(outputs) == 1 else outputs
def wrapper(*args, **kwargs):
if len(args) == 0:
raise ValueError(
'Excepted the first argument is <inputs>.')
inputs = nest.flatten(args[0])
if len(inputs) < min_num or len(inputs) > max_num:
raise ValueError('The number of <inputs> is {}, '
'not in range: [min={}, max={}].'.format(
len(inputs), min_num, max_num))
return inner_function(inputs, *args[1:], **kwargs)
def norm(input, p='fro', dim=None, keepdim=False, out=None, dtype=None):
"""Compute the norm value of elements along the given dimension.
:attr:`dim` could be negative or ``None``:
```python
x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
# A negative dimension is the last-k axis
print(torch.norm(x, dim=1))
print(torch.norm(x, dim=-1)) # Equivalent
# If ``dim`` is None, the vector-style reduction
# will be applied to return a scalar result
print(torch.norm(x)) # 9.539
# Also, ``dim`` could be a sequence of integers
print(torch.norm(x, dim=(0, 1))) # 9.539
```
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
p : {'fro', 1, 2}, optional
The norm order.
dim : Union[int, Sequence[int]], optional
The dimension to reduce.
keepdim : bool, optional, default=False
Keep the reduced dimension or not.
out : dragon.vm.torch.Tensor, optional
The output tensor.
dtype : str, optional
The data type to cast to.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
if p is None or p == 2 or p == 'fro':
op_type = 'ReduceL2'
elif p == 1:
op_type = 'ReduceL1'
else:
raise ValueError('Unsupported norm order: ' + str(p))
input = input.to(dtype=dtype)
keepdim = keepdim if dim is not None else False
dim = nest.flatten(dim) if dim is not None else dim
return Function.apply(op_type,
input.device, [input],
outputs=[out],
axes=dim,
keepdims=keepdim)
