def _arg_reduce(input, operation, dim=None, keepdim=False, top_k=1, out=None):
if dim is None: keepdim = False
dev = MakeDevice(inputs=[input])
key = '{}/{}/dim:{}/keepdim:{}/top_k:{}'.format(operation, dev, dim,
int(keepdim), top_k)
module = get_module(ArgReduce,
key,
dev,
operation=operation,
axis=dim,
keepdim=keepdim,
top_k=top_k)
return module.forward(input, out)
def _maximum(input, other, out=None):
if not isinstance(input, Tensor):
input = WrapScalar(input, 'float32', other._ctx)
dtype = other._dtype
elif not isinstance(other, Tensor):
other = WrapScalar(other, 'float32', input._ctx)
dtype = input._dtype
else:
dtype = input._dtype
ctx = MakeContext(inputs=[input])
key = 'torch/ops/maximum/{}:{}'.format(ctx[0].lower(), ctx[1])
module = get_module(Maximum, key, ctx)
return module.forward(input, other, out, dtype)
def _normal(input, shape, mean, std):
dev = MakeDevice(inputs=[input])
n_dim = len(shape)
key = 'Normal/{}/dtype:{}/n_dim:{}/mean:{}/std:{}'.format(
dev, input.dtype, n_dim, float(mean), float(std))
module = get_module(RandomNormal,
key,
dev,
n_dim=n_dim,
mean=mean,
std=std,
dtype=input.dtype)
return module.forward(input, shape)
def _uniform(input, shape, low, high):
dev = MakeDevice(inputs=[input])
n_dim = len(shape)
key = 'Uniform/{}/dtype:{}/n_dim:{}/low:{}/high:{}'.format(
dev, input.dtype, n_dim, float(low), float(high))
module = get_module(RandomUniform,
key,
dev,
n_dim=n_dim,
low=low,
high=high,
dtype=input.dtype)
return module.forward(input, shape)
def _fill(input, shape, value):
dev = MakeDevice(inputs=[input])
ndim = len(shape)
key = 'Fill/{}/dtype:{}/ndim:{}/value:{}' \
.format(dev, input.dtype, ndim, value)
module = get_module(
Fill,
key,
dev,
ndim=ndim,
value=value,
dtype=input.dtype,
)
return module.forward(input, shape)
def _assigning(output, input, starts, sizes):
if not isinstance(input, Tensor):
if isinstance(input, (tuple, list)):
input = Tensor(input, dtype=output.dtype, device=output.device)
else:
input = WrapScalar(input, output.dtype, output.device)
n_starts, n_sizes = len(starts), len(sizes)
dev = MakeDevice(inputs=[input])
key = 'Assign/{}/n_starts:{}/n_sizes:{}'.format(dev, n_starts, n_sizes)
module = get_module(Assigning,
key,
dev,
n_starts=n_starts,
n_sizes=n_sizes)
return module.forward(input, output, starts, sizes)
def _reduce(input, operation, dim=None, keepdim=False, out=None):
ctx = MakeContext(inputs=[input])
if dim is None:
dim = -1
keepdim = False
elif dim < 0:
dim = CanonicalAxis(input, dim)
key = 'torch/ops/{}/{}:{}/dim[{}]/keep_dims:{}'.format(
operation.lower(), ctx[0].lower(), ctx[1], dim, int(keepdim))
module = get_module(Reduce,
key,
ctx,
operation=operation,
axis=dim,
keep_dims=keepdim)
return module.forward(input, out)
def roi_align(feature,
rois,
pooled_h,
pooled_w,
spatial_scale,
sampling_ratio=2):
dev = MakeDevice(inputs=[feature])
key = 'RoIAlign/{}/pool_h:{}/pool_w:{}/' \
'spatial_scale:{}/sampling_ratio:{}'.format(
dev, pooled_h, pooled_w, spatial_scale, sampling_ratio)
module = get_module(RoIAlign,
key,
dev,
pooled_h=pooled_h,
pooled_w=pooled_w,
spatial_scale=spatial_scale,
sampling_ratio=sampling_ratio)
return module.forward(feature, rois)
def roi_align(feature,
rois,
pooled_h,
pooled_w,
spatial_scale,
sampling_ratio=2):
ctx = MakeContext(inputs=[feature])
key = 'torch/ops/roi_align/{}:{}/pool_h:{}/pool_w:{}/' \
'spatial_scale:{}/sampling_ratio:{}'.format(
ctx[0].lower(), ctx[1], pooled_h, pooled_w, spatial_scale, sampling_ratio)
module = get_module(RoIAlign,
key,
ctx,
pooled_h=pooled_h,
pooled_w=pooled_w,
spatial_scale=spatial_scale,
sampling_ratio=sampling_ratio)
return module.forward(feature, rois)
def one_hot(input, depth):
"""Return a ont hot tensor according to given input.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
depth : int
The depth of channels.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'OneHot/{}/depth:{}'.format(dev, depth)
module = get_module(OneHot, key, dev, depth=depth)
return module.forward(input)
def nonzero(input, out=None):
"""Return the indices of non-zero elements.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
out : dragon.vm.torch.Tensor, optional
The optional output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'NonZero/{}'.format(dev)
module = get_module(NonZero, key, dev)
return module.forward(input, out)
def unsqueeze(input, dim, out=None):
"""Returns a tensor with a dimension of size 1 inserted at the specified position.
Parameters
----------
dim : int
The dim to remove.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The new tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Unsqueeze/{}/dim:{}'.format(dev, dim if dim else 'None')
module = get_module(UnSqueeze, key, dev, dim=dim)
return module.forward(input, out=out)
def squeeze(input, dim=None, out=None):
"""Return a tensor with all the dimensions of input of size 1 removed.
Parameters
----------
dim : int
The optional dim to remove.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The new tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Squeeze/{}/dim:{}'.format(dev, dim if dim else 'None')
module = get_module(Squeeze, key, dev, dim=dim)
return module.forward(input, out=out)
def _update(
param,
grad,
op_type,
slot,
lr_mult=1.0,
decay_mult=1.0,
):
dev = MakeDevice(inputs=[param])
key = '{}/{}/{}/{}'.format(op_type, dev, slot, param.name)
module = get_module(
Update,
key,
dev,
op_type=op_type,
lr_mult=lr_mult,
decay_mult=decay_mult,
slot=slot,
)
return module.forward(param, grad)
def _resize_2d(input, op_type, dsize, fx, fy):
if dsize is None:
if fx < 0 or fy < 0:
raise ValueError('Set fx and fy if dsize is None.')
else:
if len(dsize) != 2:
raise ValueError('The dsize should be a list with 2 elements.')
if dsize is None and (fy == -1.0 or fx == -1.0):
raise RuntimeError('The dsize, fx/fy should be specified either.')
dev = MakeDevice(inputs=[input])
key = '{}/{}/dsize:{}/fx:{}/fy:{}'.format(op_type, dev,
'2' if dsize else 'none', fx, fy)
module = get_module(Resize2d,
key,
dev,
op_type=op_type,
dsize=dsize,
fx=fx,
fy=fy)
return module.forward(input, dsize)
def sqrt(input, out=None):
"""Compute the square-root of input.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Sqrt/{}'.format(dev)
module = get_module(Sqrt, key, dev)
return module.forward(input, out)
def exp(input, out=None):
"""Compute the exponential of input.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Exp/{}'.format(dev)
module = get_module(Exp, key, dev)
return module.forward(input, out)
def log(input, out=None):
"""Compute the natural logarithm of input.
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Log/{}'.format(dev)
module = get_module(Log, key, dev)
return module.forward(input, out)
def chunk(tensor, chunks, dim=0):
"""Split the input into several parts along the given axis.
Parameters
----------
tensor : dragon.vm.torch.Tensor
The input to split.
chunks : int
The number of chunks to split.
dim : int, optional, default=0
The dim to split.
Returns
-------
sequence of dragon.vm.torch.Tensor
The output chunks.
"""
dev = MakeDevice([tensor])
key = 'Chunk/{}/chunks:{}/dim:{}'.format(dev, chunks, dim)
module = get_module(Chunk, key, dev, axis=dim, chunks=chunks)
return module.forward(tensor)
def cat(seq, dim=0, out=None):
"""Concatenate the inputs along the given axis.
Parameters
----------
seq : tuple or list of vm.torch.Tensor
The sequence.
dim : int
The dim to concatenate.
out : vm.torch.Tensor or None
The optional output tensor.
Returns
-------
vm.torch.Tensor
The output tensor.
"""
ctx = MakeContext(inputs=seq, outputs=[out] if out else [])
key = 'torch/ops/cat/{}:{}/dim:{}'.format(ctx[0].lower(), ctx[1], dim)
module = get_module(Concat, key, ctx, axis=dim)
return module.forward(seq, out)
def cat(seq, dim=0, out=None):
"""Concatenate the inputs along the given axis.
Parameters
----------
seq : sequence of dragon.vm.torch.Tensor
The sequence.
dim : int, optional
The dim to concatenate.
out : dragon.vm.torch.Tensor, optional
The optional output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(seq, [out] if out else [])
key = 'Concat/{}/dim:{}'.format(dev, dim)
module = get_module(Concat, key, dev, axis=dim)
return module.forward(seq, out)
def where(condition, x, y):
"""Select elements from either ``x`` or ``y``, depending on ``condition``.
Parameters
----------
condition : dragon.vm.torch.Tensor
The byte condition tensor.
x : dragon.vm.torch.Tensor
The elements for *1*.
y : dragon.vm.torch.Tensor
The elements for *0*.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[condition, x, y])
key = 'Where/{}'.format(dev)
module = get_module(Where, key, dev)
return module.forward(condition, x, y)
def masked_select(input, mask, out=None):
"""Select the input values where mask is *1*.
Parameters
----------
input : dragon.vm.torch.Tensor
The values.
mask : dragon.vm.torch.Tensor
The mask to select values.
out : dragon.vm.torch.Tensor, optional
The optional output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice([input, mask], [out] if out else [])
key = 'MaskedSelect/{}'.format(dev)
module = get_module(MaskedSelect, key, dev)
return module.forward(input, mask, out)
def clamp(input, min=None, max=None, out=None):
"""Clamp all elements into the range [min, max].
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
min : number, optional
The min value.
max : number, optional
The max value.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Clamp/{}/min:{}/max:{}'.format(dev, min, max)
module = get_module(Clamp, key, dev, min=min, max=max)
return module.forward(input, out)
def index_select(input, dim, index, out=None):
"""Select the input values along the given axis using index.
Parameters
----------
input : dragon.vm.torch.Tensor
The values.
dim : int
The dim to gather.
index : dragon.vm.torch.Tensor
The indices.
out : dragon.vm.torch.Tensor, optional
The optional output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice([input, index], [out] if out else [])
key = 'IndexSelect/{}/dim:{}'.format(dev, dim)
module = get_module(IndexSelect, key, dev, axis=dim)
return module.forward(input, index, out)
def xw_plus_b(x, w, bias=None, transW=True, out=None):
"""Compute *matmul(x, w) + bias.*``
Parameters
----------
x : dragon.vm.torch.Tensor
The x.
w : dragon.vm.torch.Tensor
The w.
bias : dragon.vm.torch.Tensor, optional
The bias.
transW : boolean
Whether to transpose the ``w``.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[x, w] + ([bias] if bias else []))
key = 'FullyConnected/{}/transW:{}'.format(dev, transW)
module = get_module(FullyConnected, key, dev, transW=transW)
return module.forward(x, w, bias, out)
def mm(mat1, mat2, transA=False, transB=False, out=None):
"""Performs a matrix multiplication of the matrices ``mat1`` and ``mat2.``
Parameters
----------
mat1 : dragon.vm.torch.Tensor
The matrix A.
mat2 : dragon.vm.torch.Tensor
The matrix B.
transA : boolean
Whether to transpose the ``mat1``.
transB : boolean
Whether to transpose the ``mat2``.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[mat1, mat2])
key = 'Matmul/{}/transA:{}/transB:{}'.format(dev, transA, transB)
module = get_module(MM, key, dev, transA=transA, transB=transB)
return module.forward(mat1, mat2, out)
def accumulate(input, alpha=1., beta=1., out=None):
"""Compute *out = alpha * input + beta * out*
Parameters
----------
input : dragon.vm.torch.Tensor
The input tensor.
alpha : float, optional, default=1.
The value of alpha.
beta : float, optional, default=1.
The value beta.
out : dragon.vm.torch.Tensor, optional
The output tensor.
Returns
-------
dragon.vm.torch.Tensor
The output tensor.
"""
dev = MakeDevice(inputs=[input])
key = 'Accumulate/{}/alpha:{}/beta:{}'.format(dev, alpha, beta)
module = get_module(Accumulate, key, dev, alpha=alpha, beta=beta)
return module.forward(input, out)
def unsqueeze(input, dim, out=None):
ctx = MakeContext(inputs=[input])
key = 'torch/ops/unsqueeze/{}:{}/dim:{}'.format(ctx[0].lower(), ctx[1],
dim if dim else 'None')
module = get_module(UnSqueeze, key, ctx, dim=dim)
return module.forward(input, out=out)
def _index(input, starts, sizes):
nstarts, nsizes = len(starts), len(sizes)
dev = MakeDevice(inputs=[input])
key = 'Index/{}/nstarts:{}/nsizes:{}'.format(dev, nstarts, nsizes)
module = get_module(Indexing, key, dev, nstarts=nstarts, nsizes=nsizes)
return module.forward(input, starts, sizes)
