def dirac_(tensor, groups=1):
"""Fill tensor with the dirac delta function.
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
groups : number, optional, default=1
The groups of convolution.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
dimensions = tensor.ndimension()
if dimensions not in [3, 4, 5]:
raise ValueError('Only tensors with 3, 4, or 5 dimensions are supported.')
sizes = tensor.size()
if sizes[0] % groups != 0:
raise ValueError('Dimension 0 should be divisible by groups.')
out_channels_per_grp = sizes[0] // groups
min_dim = min(out_channels_per_grp, sizes[1])
with grad_mode.no_grad():
tensor.zero_()
for g in range(groups):
for d in range(min_dim):
item = [g * out_channels_per_grp + d, d]
for i in range(2, dimensions):
item.append(sizes[i] // 2)
tensor[tuple(item)] = 1
return tensor
def xavier_uniform_(tensor, gain=1):
r"""Fill tensor from a xavier uniform distribution.
.. math::
\text{tensor} \sim \mathcal{U}(-\alpha, \alpha) \\ \, \\ \,
\text{where} \quad \alpha = \text{gain} \times
\sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}}
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
gain : number, optional, default=1
The gain value.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out))
a = math.sqrt(3.0) * std
with grad_mode.no_grad():
return tensor.uniform_(-a, a)
def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
r"""Fill tensor from a kaiming uniform distribution.
.. math::
\text{tensor} \sim \mathcal{U}(-\alpha, \alpha) \\ \, \\ \,
\text{where} \quad \alpha = \text{gain} \times
\sqrt{\frac{3}{\text{fan}}}
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
a : number, optional, default=0
The negative slope to compute gain value.
mode : {'fan_in', 'fan_out'}, optional
The mode to compute fans.
nonlinearity : str, optional, default='leaky_relu'
The nonlinearity to compute gain value.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
fan = _calculate_correct_fan(tensor, mode)
gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan)
bound = math.sqrt(3.0) * std
with grad_mode.no_grad():
return tensor.uniform_(-bound, bound)
def xavier_normal_(tensor, gain=1):
r"""Fill tensor from a xavier normal distribution.
.. math::
\text{tensor} \sim \mathcal{N}(0, \sigma^{2}) \\ \, \\ \,
\text{where} \quad \sigma = \text{gain} \times
\sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}}
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
gain : number, optional, default=1
The gain value.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
std = gain * math.sqrt(2.0 / (fan_in + fan_out))
with grad_mode.no_grad():
return tensor.normal_(0, std)
def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
r"""Fill tensor from a kaiming normal distribution.
.. math::
\text{tensor} \sim \mathcal{N}(0, \sigma^{2}) \\ \, \\ \,
\text{where} \quad \sigma = \text{gain} \times
\sqrt{\frac{1}{\text{fan}}}
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
a : number, optional, default=0
The negative slope to compute gain value.
mode : {'fan_in', 'fan_out'}, optional
The mode to compute fans.
nonlinearity : str, optional, default='leaky_relu'
The nonlinearity to compute gain value.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
fan = _calculate_correct_fan(tensor, mode)
gain = calculate_gain(nonlinearity, a)
std = gain / math.sqrt(fan)
with grad_mode.no_grad():
return tensor.normal_(0, std)
def _apply(self, fn):
for module in self.children():
module._apply(fn)
for param in self._parameters.values():
if param is not None:
with grad_mode.no_grad():
fn(param)
for key, buf in self._buffers.items():
if buf is not None:
self._buffers[key] = fn(buf)
return self
def apply(function, *args, **kwargs):
"""Apply function and create a checkpoint."""
kwargs.pop('preserve_rng_state', True)
variable_scope = kwargs.pop('variable_scope', 'Buffer')
original_variable_scope = context.get_variable_scope(True)
if kwargs:
raise ValueError('Unexpected keyword arguments: ' +
','.join(arg for arg in kwargs))
# Run function.
graph_tape = tapes.Tape()
graph_tape._tracing = True # Enable tracing.
graph_tape._checkpointing = True # Enable checkpointing.
graph_tape._original_variable_scope = original_variable_scope
with grad_mode.no_grad(), graph_tape:
with context.variable_scope(variable_scope):
outputs = function(*args)
# Collect involving tensors.
tensor_inputs, tensor_outputs = [], []
for arg in args:
if isinstance(arg, Tensor):
tensor_inputs.append(arg)
for arg in nest.flatten(outputs):
if isinstance(arg, Tensor):
tensor_outputs.append(arg)
# Fill tape with function context.
op_tape = tapes.OrderedTape()
op_handle = workspace.get_workspace().create_handle('Checkpoint')
op_tape.add_element(proto_util.make_operator_def(
op_type='Checkpoint',
name=op_handle,
inputs=[input.id for input in tensor_inputs],
outputs=[output.id for output in tensor_outputs],
defs=[v.SerializeAs() for v in graph_tape.get_elements()],
buffer_scope=variable_scope,
to_impl=True))
op_tape.add_handle(op_handle)
op_tape.merge_handles(graph_tape.get_handles())
# Save input tensors for backward.
for input in tensor_inputs + graph_tape.get_sources():
op_tape.add_source(input)
# Save tape for backward.
for output in tensor_outputs:
output._tape = op_tape
output._requires_grad = True
return outputs
def zeros_(tensor):
r"""Fill tensor with zeros.
.. math:: \text{tensor} \leftarrow 0
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
with grad_mode.no_grad():
return tensor.fill_(0)
def constant_(tensor, val):
r"""Fill tensor with the scalar value.
.. math:: \text{tensor} \leftarrow \text{value}
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
val : number
The value to fill.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
with grad_mode.no_grad():
return tensor.fill_(val)
def eye_(tensor):
r"""Fill tensor as the identity matrix.
.. math:: \text{tensor} \leftarrow \text{diag}(1, 1, ..., 1)
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
if tensor.ndimension() != 2:
raise ValueError('Only tensors with 2 dimensions are supported.')
with grad_mode.no_grad():
init_funcs.eye(*tensor.shape, out=tensor, requires_grad=tensor.requires_grad)
return tensor
def uniform_(tensor, a=0, b=1):
r"""Fill tensor from an uniform distribution.
.. math:: \text{tensor} \sim \mathcal{U}(\alpha, \beta)
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
a : number, optional, default=-1
The value to :math:`\alpha`.
b : number, optional, default=1
The value to :math:`\beta`.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
with grad_mode.no_grad():
return tensor.uniform_(a, b)
def normal_(tensor, mean=0, std=1):
r"""Fill tensor from a normal distribution.
.. math:: \text{tensor} \sim \mathcal{N}(\mu, \sigma^{2})
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input tensor.
mean : number, optional, default=0
The value to :math:`\mu`.
std : number, optional, default=1
The value to :math:`\sigma`.
Returns
-------
dragon.vm.torch.Tensor
The input tensor.
"""
with grad_mode.no_grad():
return tensor.normal_(mean, std)
