megengine.functional

import megengine.functional as F

注解

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

参见

用户指南中对于 使用 Functional 操作与计算 有另外一套分类逻辑,可作为参考。

全局tensor操作

注解

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

Creation Functions

arange

返回半开区间 [start, stop) 内均匀间隔的值组成的1维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

返回一个特定形状的元素全为 0 的张量。

zeros_like

返回一个跟输入tensor拥有相同shape和dtype的全零tensor。

ones

返回一个新创建的拥有指定shape且被1填充的tensor。

ones_like

使用输入 Tensor 来填充输出 Tensor,保持相同的形状和数据类型。

full

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

full_like

返回一个与输入张量形状相同的且值为给定值的零张量

Manipulation Functions

copy

把张量复制到另一个设备上。

reshape

在不改变数据的情况下更改 Tensor 的形状。

flatten

通过将子张量从 start_axis 维展平到 end_axis 维,实现对张量的重塑(reshape)。

transpose

根据给定模板交换形状和步长。

swapaxes

交换一个 tensor 的两个坐标轴。

broadcast_to

将张量广播至指定的形状。

expand_dims

在给定轴之前添加维度。

squeeze

移除形状中下标为1的维度。

concat

拼接一些张量

stack

沿新轴对张量进行拼接。

split

将输入张量拆分成更小的张量。

tile

通过重复 inp 数据 reps 次来构造一个数组。

repeat

重复数组的元素。

roll

沿给定轴axis循环滚动张量

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

返回在给定轴 axisinp 张量每一行进行 \(L_p\) 归一化后的结果。

参见

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

Linear Algebra Functions

dot

计算两个向量 inp1inp2 的点积。

matinv

计算一批矩阵的逆;输入必须满足 [...,n,n] 的形状。

matmul

对矩阵 inp1inp2 进行矩阵乘法。

svd

计算矩阵(或一组矩阵) inp 的奇异值分解。

Indexing Functions

gather

使用 indexaxis 上的 inp 聚合数据。

scatter

source 张量中所有的值通过 index 张量指定的索引位置上写入到输入张量中。

cond_take

如果在满足特定条件,则从数据中获取元素。

where

根据 mask 条件选择张量 xy 作为对应的输出。

Searching Functions

argmin

返回沿给定轴的最小值的索引。

argmax

返回沿给定轴的最大值的索引。

Sorting Functions

sort

返回经过排序后的张量。

argsort

返回经过排序后的张量。

topk

按行排序,选出二维矩阵中 Top-K (默认情况下)个最小元素。

Metric functions

topk_accuracy

根据给定的预测的logits和真实值标签计算分类准确率。

神经网络运算

卷积运算

conv1d

一维卷积运算。

conv2d

二维卷积运算。

conv3d

三维卷积运算。

local_conv2d

对4D输入tensor进行通道分组并应用非统一kernel的空间卷积操作。

conv_transpose2d

二维转置卷积运算。

conv_transpose3d

3D transposed convolution operation.

deformable_conv2d

可变形卷积

sliding_window

从多batch输入tensor中提取滑窗大小的局部分块。

sliding_window_transpose

对滑动窗口对应的输入数据求和。

池化函数

avg_pool2d

对输入进行二维平均池化。

max_pool2d

对输入张量进行二维最大池化。

adaptive_avg_pool2d

对输入进行二维平均池化。

adaptive_max_pool2d

对输入数据进行2D最大池化。

deformable_psroi_pooling

可变形的位置敏感的感兴趣区域池化

非线性激活函数

sigmoid

逐元素计算 1 / ( 1 + exp( -x ) ).

hsigmoid

逐元素计算 relu6(x + 3) / 6.

hswish

逐元素计算 x * relu6(x + 3) / 6.

relu

逐元素计算 max(x, 0).

relu6

逐元素计算 min(max(x, 0), 6).

prelu

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

leaky_relu

LeakyReLU 函数。

silu

逐元素应用 SiLU 函数: x * sigmoid(x)

gelu

逐元素应用函数:

softplus

逐元素应用函数:

softmax

应用 \(\text{softmax}(x)\) 函数

logsoftmax

\(\log(\text{softmax}(x))\) 函数应用于一个 n 维 Tensor 输入上

logsigmoid

逐元素应用函数:

logsumexp

计算输入数据沿给定维度 axis 的指数之和的对数。该计算是数值稳定的。

标准化函数

batch_norm

对输入进行批标准化。

sync_batch_norm

对输入进同步批标准化。

layer_norm

Applies layer normalization to the input.

local_response_norm

对输入 Tensor 进行局部响应归一化。

线性函数

linear

对输入张量进行线性变换。

Dropout functions

dropout

返回一个新的张量,其中每个元素以 P = drop_prob 的概率随机设置为零。

Sparse functions

one_hot

对输入张量进行 one-hot 编码。

indexing_one_hot

对一些轴进行 One-hot 索引。

embedding

应用查找表进行 embedding。

Loss functions

l1_loss

计算预测值 \(x\) 和标签值 \(y\) 的每个元素之间的平均绝对误差(MAE)。

square_loss

计算预测值 \(x\) 和标签值 \(y\) 之间的均方误差(平方L2范数)。

hinge_loss

计算支持向量机 SVM 中经常使用的 hinge loss。

binary_cross_entropy

计算 binary cross entropy loss(默认使用 logits)。

cross_entropy

计算 multi-class cross entropy loss(默认使用 logits)。

ctc_loss

计算 Connectionist Temporal Classification loss 。

视觉运算

cvt_color

将图像从一种格式转化成另一种格式。

pad

填充输入张量。

pixel_shuffle

将形状为 (..., C * r^2, H, W) 的张量的元素重排为一个形状为 (..., C, H * r, W * r) 的张量,其中 r 是一个放大系数, ... 是零个或更多的 batch 维度。

interpolate

根据给定的大小或缩放因子将输入张量进行上/下采样。

remap

对按批组织的二维图像进行重映射变换。

warp_affine

Batched affine transformation on 2D images.

warp_perspective

对按批组织的二维图像进行透视变换。

roi_pooling

在指定输入的感兴趣区域上执行池化,并获得固定大小的特征图, 如Faster RCNN中所述

roi_align

RoI Align 是在指定输入的感兴趣区域上获得特征图,如 Mask R-CNN 论文中所述。

nms

根据小方框的 IoU 值进行 非最大值抑制(NMS)。

correlation

对输入进行相关性计算。

nvof

NVIDIA光流SDK的实现

Quantized functions

conv_bias_activation

带有激活函数的卷积偏置操作,仅用于推理

batch_conv_bias_activation

带有激活函数的批量卷积偏置,仅用于推理

conv_transpose2d

rtype

Tensor

分布式算子

all_gather

在指定的组中收集张量,并沿第一维对其进行拼接。

all_reduce_max

在指定组中以 max 操作来对张量进行规约操作。

all_reduce_min

在指定组中以 min 操作来对张量进行规约操作。

all_reduce_sum

在指定组中以求和操作来对张量进行规约操作。

all_to_all

每个进程将输入张量分散到所有进程,并返回收集的张量。

broadcast

从根进程向其他进程广播张量数据。

reduce_scatter_sum

通过求和规约指定组中的张量,并在第一维度将其拆分。

collective_comm

用于集合通信的辅助函数。

remote_send

发送张量到另一个进程。

remote_recv

从另一个进程接收张量。

reduce_sum

在指定组中以 sum操作来对张量进行规约操作。

gather

在指定的组中收集张量。

scatter

在根进程中沿第一维度切分张量。

External functions

extern_opr_subgraph

加载序列化的extern opr子图并假执行运算符。

tensorrt_runtime_opr

cambricon_runtime_opr

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

atlas_runtime_opr

加载 Atlas 模型。

magicmind_runtime_opr

加载序列化的MagicMind模型作为MegEngine中的运行时操作符。