megengine.functional#

import megengine.functional as F

Note

顾名思义,megengine.functional 模块中包含着所有与 Tensor 有关的计算接口:

See also

用户指南中对于 Use Functional operations and calculations 有另外一套分类逻辑,可作为参考。

General tensor operations#

Note

该部分的 API 设计接纳了 Python 数据 API 标准联盟 中的倡导, 尽可能地向 NumPy API 靠拢。

Creation Functions#

arange

Returns evenly spaced values within the half-open interval [start, stop) as a one-dimensional tensor.

linspace

Returns evenly spaced numbers over a specified interval.

eye

Returns a two-dimensional tensor with ones on the diagonal and zeros elsewhere.

diag

Extract a diagonal or construct a diagonal tensor.

zeros

Returns a new tensor having a specified shape and filled with zeros.

zeros_like

Returns a tensor filled with zeros with the same shape and data type as input tensor.

ones

Returns a new tensor having a specified shape and filled with ones.

ones_like

Returns a tensor filled with ones with the same shape and data type as input tensor.

full

Returns a new tensor having a specified shape and filled with given value.

full_like

Returns a tensor filled with given value with the same shape as input tensor.

Manipulation Functions#

copy

Copies tensor to another device.

reshape

Reshapes a tensor without changing its data.

flatten

Reshapes the tensor by flattening the sub-tensor from dimension start_axis to dimension end_axis.

transpose

Swaps shapes and strides according to given pattern.

swapaxes

Interchange two axes of a tensor.

broadcast_to

Broadcasts a tensor to given shape.

expand_dims

Adds dimension before given axis.

squeeze

Removes dimension of shape 1.

concat

Concat some tensors

stack

Concats a sequence of tensors along a new axis.

split

Splits the input tensor into several smaller tensors.

tile

Construct an array by repeating inp the number of times given by reps.

repeat

Repeat elements of an array.

roll

Roll the tensor along the given axis(or axes).

Arithmetic operations#

add

Element-wise addition.

sub

Element-wise subtraction.

mul

Element-wise multiplication.

div

Element-wise division.

floor_div

Element-wise floor division.

neg

Element-wise negation.

abs

Element-wise \(\operatorname{abs}(x)\) function.

sign

Element-wise sign.

pow

Element-wise power.

mod

Element-wise \(\operatorname{mod}(x, y)\) function.

sqrt

Element-wise \(\operatorname{sqrt}(x)\) function.

square

Element-wise \(x^2\) function.

maximum

Element-wise maximum of tensor elements.

minimum

Element-wise minimum of tensor elements.

round

Element-wise \(\operatorname{round}(x)\) function.

ceil

Element-wise \(\lceil x \rceil\) function.

floor

Element-wise \(\lfloor x \rfloor\) function.

clip

Element-wise clipping function.

exp

Element-wise \(e^x\) function.

expm1

Element-wise \(e^x - 1\) function.

log

Element-wise \(\operatorname{log}(x)\) function.

log1p

Element-wise \(\log(1 + x)\) function.

logaddexp

Element-wise \(\log(e^x + e^y)\) function.

Trigonometric functions#

sin

Element-wise \(\sin(x)\) function.

cos

Element-wise \(\cos(x)\) function.

tan

Element-wise \(\tan(x)\) function.

asin

Element-wise \(\arcsin(x)\) function.

acos

Element-wise \(\arccos(x)\) function.

atan

Element-wise \(\arctan(x)\) function.

atan2

Element-wise \(\arctan(\frac{y}{x})\) function.

Hyperbolic functions#

sinh

Element-wise \(\sinh(x)\) function.

cosh

Element-wise \(\cosh(x)\) function.

tanh

Element-wise \(\tanh(x)\) function.

acosh

Element-wise \(\cosh^{-1}(x)\) function.

asinh

Element-wise \(\sinh^{-1}(x)\) function.

atanh

Element-wise \(\tanh^{-1}(x)\) function.

Bit operations#

left_shift

Element-wise left shift.

right_shift

Element-wise right shift.

Logic functions#

isnan

Element-wise NaN check.

isinf

Element-wise infinity check.

logical_and

Element-wise logical AND.

logical_not

Element-wise logical NOT.

logical_or

Element-wise logical OR.

logical_xor

Element-wise logical XOR.

greater

Element-wise greater-than comparison.

greater_equal

Element-wise greater-than-or-equal-to comparison.

less

Element-wise less-than comparison.

less_equal

Element-wise less-than-or-equal-to comparison.

equal

Element-wise equality comparison.

not_equal

Element-wise inequality comparison.

Statistical Functions#

sum

Calculates the sum of tensor elements over a given axis (or axes).

cumsum

Calculates the cumulative sum of tensor elements over a given axis.

prod

Calculates the product of tensor elements over a given axis (or axes).

mean

Calculates the mean of tensor elements over a given axis (or axes).

min

Calculates the minimum of tensor elements over a given axis (or axes).

max

Calculates the maximum of tensor elements over a given axis (or axes).

var

Calculates the variance of tensor elements over a given axis (or axes).

std

Calculates the standard deviation of tensor elements over a given axis (or axes).

norm

Calculates the norm of tensor elements over a given axis.

normalize

Performs \(L_p\) normalization of input tensor along given axis.

See also

想要返回 min, max 的索引而不是元素值,请参考 Searching Functions

Linear Algebra Functions#

dot

Computes dot-product of two vectors inp1 and inp2.

matinv

Computes the inverse of a batch of matrices; input must has shape [..., n, n].

matmul

Performs a matrix multiplication of the matrices inp1 and inp2.

svd

Computes the singular value decomposition of a matrix (or a stack of matrices) inp.

Indexing Functions#

gather

Gathers data from input tensor on axis using index.

scatter

Writes all values from the tensor source into input tensor at the indices specified in the index tensor.

cond_take

Takes elements from data if specific condition is satisfied on mask.

where

Selects elements either from Tensor x or Tensor y, according to mask.

Searching Functions#

argmin

Returns the indices of the minimum values along given axis.

argmax

Returns the indices of the maximum values along given axis.

Sorting Functions#

sort

Returns sorted tensor and the indices would sort the input tensor.

argsort

Returns the indices that would sort the input tensor.

topk

Selects the Top-K (by default) smallest elements of 2d matrix by row.

Metric functions#

topk_accuracy

Calculates the classification accuracy given predicted logits and ground-truth labels.

Neural network operations#

Convolution functions#

conv1d

1D convolution operation.

conv2d

2D convolution operation.

conv3d

3D convolution operation.

local_conv2d

Applies a spatial convolution with untied kernels over an groupped channeled input 4D tensor.

conv_transpose2d

2D transposed convolution operation.

conv_transpose3d

3D transposed convolution operation.

deformable_conv2d

Deformable Convolution.

sliding_window

Extracts sliding local blocks from a batched input tensor.

sliding_window_transpose

Sum over the sliding windows on the corresponding input location.

Pooling functions#

avg_pool2d

Applies 2D average pooling over an input tensor.

max_pool2d

Applies a 2D max pooling over an input tensor.

adaptive_avg_pool2d

Applies a 2D average adaptive pooling over an input.

adaptive_max_pool2d

Applies a 2D max adaptive pooling over an input.

deformable_psroi_pooling

Deformable PSROI(Position Sensitive Region of Interest) Pooling.

Non-linear activation functions#

sigmoid

Element-wise 1 / ( 1 + exp( -x ) ).

hsigmoid

Element-wise relu6(x + 3) / 6.

hswish

Element-wise x * relu6(x + 3) / 6.

relu

Element-wise max(x, 0).

relu6

Element-wise min(max(x, 0), 6).

prelu

Element-wise max(x, 0) + y * min(x, 0).

leaky_relu

Element-wise LeakyReLU function

silu

Applies the element-wise Sigmoid Linear Unit function, i.e. x * sigmoid(x).

gelu

Applies the element-wise function:

softplus

Applies the element-wise function:

softmax

Applies a \(\text{softmax}(x)\) function.

logsoftmax

Applies the \(\log(\text{softmax}(x))\) function to an n-dimensional input tensor.

logsigmoid

Applies the element-wise function:

logsumexp

Calculates the logarithm of the inputs' exponential sum along the given axis.

Normalization functions#

batch_norm

Applies batch normalization to the input.

sync_batch_norm

Applies synchronized batch normalization to the input.

layer_norm

Applies layer normalization to the input.

local_response_norm

Apply local response normalization to the input tensor.

Linear functions#

linear

Applies a linear transformation to the input tensor.

multi-head attention functions#

Dropout functions#

dropout

Returns a new tensor where each of the elements are randomly set to zero with probability P = drop_prob.

Sparse functions#

one_hot

Performs one-hot encoding for the input tensor.

indexing_one_hot

One-hot indexing for some axes.

embedding

Applies lookup table for embedding.

Loss functions#

l1_loss

Calculates the mean absolute error (MAE) between each element in the pred \(x\) and label \(y\).

square_loss

Calculates the mean squared error (squared L2 norm) between each element in the pred \(x\) and label \(y\).

hinge_loss

Caculates the hinge loss which is often used in SVM.

binary_cross_entropy

Computes the binary cross entropy loss (using logits by default).

cross_entropy

Computes the multi-class cross entropy loss (using logits by default).

ctc_loss

The Connectionist Temporal Classification loss.

Vision functions#

cvt_color

Convert images from one format to another

pad

Pads the input tensor.

pixel_shuffle

Rearranges elements in a tensor of shape (..., C * r^2, H, W) to a tensor of shape (..., C, H * r, W * r), where r is an upscale factor, where ... is zero or more batch dimensions.

interpolate

Down/up samples the input tensor to either the given size or with the given scale_factor.

remap

Applies remap transformation to batched 2D images.

warp_affine

Batched affine transformation on 2D images.

warp_perspective

Applies perspective transformation to batched 2D images.

roi_pooling

Applies RoI (Region of Interest) pooling on input feature, as described in Faster RCNN.

roi_align

Applies RoI (Region of Interest) align on input feature, as described in Mask R-CNN.

nms

Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union(IoU).

correlation

Applies correlation to inputs.

nvof

Implements NVIDIA Optical Flow SDK.

Quantized functions#

conv_bias_activation

Convolution bias with activation operation, only for inference.

batch_conv_bias_activation

Batch convolution bias with activation operation, only for inference.

conv_transpose2d

rtype:

Tensor

Distributed functions#

all_gather

Gather tensors across the specified group and concat them at first dimension.

all_reduce_max

Reduce tensors with max operation on each value across the specified group.

all_reduce_min

Reduce tensors with min operation on each value across the specified group.

all_reduce_sum

Reduce tensors with sum operation on each value across the specified group.

all_to_all

Each process scatter input tensor to all processes and return gathered tensor.

broadcast

Broadcast tensor data from root process to others.

reduce_scatter_sum

Reduce tensors across the specified group by sum and split them at first dimension.

collective_comm

Helper function for applying collective communication functions.

remote_send

Send tensor to another process.

remote_recv

Receive a tensor from another process.

reduce_sum

Reduce tensor data across the specified group by sum.

gather

Gather tensors across the specified group.

scatter

Split tensor in root process at first dimension.

External functions#

extern_opr_subgraph

Load a serialized extern opr subgraph and fake execute the operator.

tensorrt_runtime_opr

cambricon_runtime_opr

Load a serialized Cambricon model as a runtime operator in MegEngine.

atlas_runtime_opr

Load a serialized Atlas model as a runtime operator in MegEngine.

magicmind_runtime_opr

Load a serialized MagicMind model as a runtime operator in MegEngine.