def multinomial(input, num_samples, eps=0., out=None): """Return a tensor where each row contains ``num_samples``, sampled from the multinomial distribution. Parameters ---------- input : dragon.vm.torch.Tensor The input tensor. num_samples : int The number of samples. eps : float, optional, default=0. The prob to a uniform sampling. Returns ------- dragon.vm.torch.Tensor The output tensor. """ dev = MakeDevice(inputs=[input]) key = 'Multinomial/{}' \ '/num_samples:{}' \ '/eps:{}'.format(dev, num_samples, eps) module = get_module( Multinomial, key, dev, eps=eps, num_samples=num_samples, ) return module.forward(input, out)
def multinomial(input, num_samples, normalize=False, out=None): """Return a tensor where each row contains ``num_samples``, sampled from the multinomial distribution. Parameters ---------- input : dragon.vm.torch.Tensor The input tensor. num_samples : int The number of samples. normalize : boolean, optional, default=False Whether to normalize the inputs. Returns ------- dragon.vm.torch.Tensor The output tensor. """ dev = MakeDevice(inputs=[input]) key = 'Multinomial/{}/num_samples:{}/normalize:{}'.format( dev, num_samples, normalize) module = get_module(Multinomial, key, dev, num_samples=num_samples, normalize=normalize) return module.forward(input, out)
def _type_to(input, dtype='float32', inplace=False): if dtype == input.dtype: return input dev = MakeDevice(inputs=[input]) key = 'Cast/{}/dtype:{}/inplace:{}'.format( dev, dtype, 'true' if inplace else 'false') module = get_module(Cast, key, dev, dtype=dtype, inplace=inplace) return module.forward(input)
def channel_shuffle(input, dim=0, group=1, out=None): """Shuffle channels between groups along the given axis. Parameters ---------- input : dragon.vm.torch.Tensor The input tensor. dim : int, optional, default=0 The axis of channels. group : int, optional, default=1 The number of groups. out : dragon.vm.torch.Tensor, optional The output tensor. Returns ------- dragon.vm.torch.Tensor The new tensor. """ dev = MakeDevice([input]) key = 'ChannelShuffle/{}/dim:{}/group:{}'.format(dev, dim, group) module = get_module( ChannelShuffle, key, dev, axis=dim, group=group, ) return module.forward(input, out)
def _allreduce(grads): if not isinstance(grads, (list, tuple)): grads = [grads] dev = MakeDevice(inputs=grads) mode = mpi.GetParallelMode() + '_ALLREDUCE' key = 'Collective/{}/{}'.format(dev, mode.lower()) module = get_module(Collective, key, dev, mode=mode) return module.forward(grads)
def gather(input, dim, index, out=None): """Gather the input values along the given axis. Note that it is a tensorflow style gather, which takes a vector index, values of other dimension will be copied automatically. 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(inputs=[input, index], outputs=[out] if out else []) key = 'Gather/{}/dim:{}'.format(dev, dim) module = get_module(Gather, key, dev, axis=dim) return module.forward(input, index, out)
def roi_pool( feature, rois, pooled_h, pooled_w, spatial_scale, ): dev = MakeDevice(inputs=[feature]) key = 'RoIPool/{}' \ '/pool_h:{}' \ '/pool_w:{}' \ '/spatial_scale:{}' \ .format(dev, pooled_h, pooled_w, spatial_scale) module = get_module( RoIPool, key, dev, pooled_h=pooled_h, pooled_w=pooled_w, spatial_scale=spatial_scale, ) return module.forward(feature, rois)
def _rfundamental(input, value, op='RAdd', out=None): if not isinstance(value, Tensor): value = WrapScalar(value, input.dtype, input.device) dev = MakeDevice(inputs=[input, value]) key = '{}/{}'.format(op, dev) module = get_module(Fundamental, key, dev, op_type=op) return module.forward(value, 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 minimum(input, other, out=None): """Return the min value of given two tensors. Parameters ---------- input : dragon.vm.torch.Tensor or number The input tensor. other : dragon.vm.torch.Tensor or number The input tensor. out : dragon.vm.torch.Tensor, optional The output tensor. Returns ------- dragon.vm.torch.Tensor The output tensor. """ if not isinstance(input, Tensor): input = WrapScalar(input, other.dtype, other.device) elif not isinstance(other, Tensor): other = WrapScalar(other, input.dtype, input.device) dev = MakeDevice(inputs=[input]) key = 'Minimum/{}'.format(dev) module = get_module(Minimum, key, dev) return module.forward(input, other, out)
def _reshape(input, shape, shape_like=None): if shape_like is not None: shape = shape_like.shape dev = MakeDevice(inputs=[input]) ndim = len(shape) key = 'Reshape/{}/ndim:{}'.format(dev, ndim) module = get_module(Reshape, key, dev, ndim=ndim) return module.forward(input, shape)
def _compare(input, other, operation, out=None): if not isinstance(other, Tensor): other = WrapScalar(other, input.dtype, input.device) dev = MakeDevice(inputs=[input, other]) key = 'Compare/{}/{}'.format(operation, dev) module = get_module(Compare, key, dev, operation=operation) return module.forward(input, other, out)
def _accumulate(grads): if len(grads) == 0: return if not isinstance(grads, (list, tuple)): grads = [grads] dev = MakeDevice(inputs=grads) key = 'Accumulate/{}/alpha:1./beta:1.'.format(dev) module = get_module(_Accumulate, key, dev) return module.forward(grads)
def _masked_assign(output, mask, input): 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) dev = MakeDevice(inputs=[input]) key = 'MaskedAssign/{}'.format(dev) module = get_module(MaskedAssign, key, dev) return module.forward(input, output, mask)
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 _reduce(input, operation, dim=None, keepdim=False, out=None): if dim is None: keepdim = False dev = MakeDevice(inputs=[input]) key = '{}/{}/dim:{}/keepdim:{}'.format(operation, dev, dim, int(keepdim)) module = get_module(Reduce, key, dev, operation=operation, dim=dim, keepdim=keepdim) return module.forward(input, out)
def _assign(output, starts, sizes, input): 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) nstarts, nsizes = len(starts), len(sizes) dev = MakeDevice(inputs=[input]) key = 'Assign/{}/nstarts:{}/nsizes:{}'.format(dev, nstarts, nsizes) module = get_module(Assign, key, dev, nstarts=nstarts, nsizes=nsizes) return module.forward(input, output, starts, sizes)
def _fill(input, shape, value): dev = MakeDevice(inputs=[input]) n_dim = len(shape) key = 'Fill/{}/dtype:{}/n_dim:{}/value:{}'.format(dev, input.dtype, n_dim, value) module = get_module(Fill, key, dev, n_dim=n_dim, value=value, 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 _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 _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 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 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 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 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 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 _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 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)