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

返回指定间隔的等距数列。

eye

返回一个二维张量,其对角线上值均为1,其他位置值为0。

diag

如果 inp 是一维tensor,返回值是一个以 inp 为对角线的二维tensor。

zeros

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

zeros_like

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

ones

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

ones_like

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

full

返回一个形状为shape类型为type元素值为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

计算输入 \(x\) 与输入 \(y\) 的逐元素和。

sub

Calculates the difference for each element \(x_i\) of the input tensor \(x\) with the respective element \(y\) of the input tensor \(y\).

mul

计算输入张量 x 的每个元素 \(x_i\) 与输入张量 \(y\) 的相应元素 \(y_i\) 的乘积。

div

计算输入 \(x\) 与输入 \(y\) 的逐元素除。

floor_div

将输入张量 \(x\) 的每个元素 \(x_i\) 除以输入张量 \(y_i\) 的相应元素 \(y\) 的结果四舍五入到最大 (即最接近到 +infinity) 不大于除法结果的整数值数。

neg

计算输入张量 \(x\) 的每个元素 \(x_i\) 的负数 (即 \(y_i = -x_i\))。

abs

逐元素取绝对值。

sign

返回一个新的张量,该张量表明每个元素的符号。

pow

一种特定实现的近似取幂运算,将张量 \(x\) 中的每一个元素 \(x_i\) (底数) 跟张量 \(y\) 中每一个元素 \(y_i\) (指数)做幂运算。

mod

返回输入张量 x 的每个元素 x_i 和输入张量 y 的相应元素 y_i 的除法余数。

sqrt

逐元素取平方根。

square

逐元素取平方。

maximum

逐元素的最大数组元素。

minimum

逐元素的最小数组元素。

round

逐元素四舍五入到整数。

ceil

逐元素上取整。

floor

逐元素计算下取整函数。

clip

将输入张量的所有元素限制在 [ lower, upper ] 区间并返回相应的结果:

exp

逐元素计算 ex 次方。

expm1

逐元素计算 expm1 函数。

log

逐元素计算以 e 为底的对数。

log1p

逐元素计算 (x+1)e 为底的对数。

logaddexp

逐元素数值稳定的 log(exp(x) + exp(y))

Trigonometric functions

sin

逐元素计算正弦函数。

cos

逐元素计算余弦函数。

tan

逐元素计算正切函数。

asin

逐元素计算反正弦函数。

acos

逐元素计算反余弦函数。

atan

逐元素计算反正切函数。

atan2

逐元素计算两个参数的反正切函数。

Hyperbolic functions

sinh

逐元素计算双曲正弦函数。

cosh

逐元素计算双曲余弦函数。

tanh

逐元素计算双曲正切值。

acosh

逐元素计算反双曲余弦函数。

asinh

逐元素计算反双曲正弦函数。

atanh

逐元素计算反双曲正切函数。

Bit operations

left_shift

逐元素计算 x << y。

right_shift

逐元素计算 x >> y。

Logic functions

isnan

返回一个新的张量,该张量表明每个元素是否为非数值类型 NaN.

isinf

返回一个张量,它表示每个元素是否是无穷大值 Inf.

logical_and

逐元素进行逻辑与运算。

logical_not

逐元素进行逻辑非运算。

logical_or

逐元素进行逻辑或运算。

logical_xor

逐元素进行逻辑异或运算。

greater

逐元素计算 x 大于 y.

greater_equal

逐元素判断 x 大于或者等于 y.

less

逐元素的计算 x 小于 y.

less_equal

逐元素计算 x 小于或等于 y.

equal

逐元素判断是否相等。

not_equal

逐元素计算 x 不等于 y.

Statistical Functions

sum

返回输入张量在给定轴上所有元素求和后的结果。

cumsum

返回输入张量在给定轴上所有元素求累加和的结果

prod

返回输入张量在给定轴上所有元素求乘积后的结果。

mean

返回输入张量在给定轴上所有元素求平均后的结果。

min

返回输入张量在给定轴上所有元素求最小值后的结果。

max

返回输入张量在给定轴上所有元素求最大值后的结果。

var

返回输入张量在给定轴上所有元素求方差后的结果。

std

返回输入张量在给定轴上所有元素求标准差后的结果。

norm

计算在给定轴 axisinp 张量每一行的 p-范数。

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 PReLU function.

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中的运行时操作符。