megengine.functional#

import megengine.functional as F

备注

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

参见

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

全局tensor操作#

备注

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

创建功能#

arange

返回半开区间 [start, stop) 内均匀间隔的值组成的1维Tensor。

linspace

返回指定区间内均匀分布的数字。

eye

返回一个对角线全1而其他元素全0的二维张量。

diag

提取或者构造一个对角线向量

zeros

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

zeros_like

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

ones

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

ones_like

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

full

返回一个新的拥有指定 shape 且全部元素均为指定值的张量。

full_like

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

操作功能#

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

逐元素相加。

sub

逐元素减法。

mul

逐元素乘法

div

逐元素除法。

floor_div

逐元素取整除法

neg

逐元素取反。

abs

将函数 \(\operatorname{abs}(x)\) 应用于每一个元素。

sign

元素求 sign

pow

逐元素幂。

mod

逐元素 \(\operatorname{mod}(x, y)\) 函数

sqrt

元素 \(\operatorname{sqrt}(x)\) 函数

square

元素 \(x^2\) 函数

maximum

Tensor中的逐元素最大值。

minimum

Tensor中的逐元素最小值。

round

逐元素的 \(\operatorname{round}(x)\) 函数.

ceil

逐元素 \(\lceil x \rceil\) 函数。

floor

逐元素 \(\lfloor x \rfloor\) 操作。

clip

逐元素裁切函数。

exp

逐元素 \(e^x\) 操作。

expm1

逐元素 \(e^x - 1\) 操作。

log

Element-wise \(\operatorname{log}(x)\) 方法。

log1p

Element-wise \(\log(1 + x)\) 方法.

logaddexp

Element-wise \(\log(e^x + e^y)\) 方法。

三角函数#

sin

元素 \(\sin(x)\) 函数

cos

逐个元素的 \(\cos(x)\) 函数。

tan

元素 \(\tan(x)\) 函数

asin

逐元素 \(\arcsin(x)\) 函数。

acos

将函数 \(\arccos(x)\) 应用于每一个元素

atan

逐元素 \(\arctan(x)\) 函数。

atan2

逐元素 \(\arctan(\frac{y}{x})\) 函数。

双曲函数#

sinh

元素 \(\sinh(x)\) 函数

cosh

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

tanh

元素 \(\tanh(x)\) 函数

acosh

将函数 \(\cosh^{-1}(x)\) 应用于每一个元素。

asinh

逐元素 \(\sinh^{-1}(x)\) 函数。

atanh

逐元素 \(\tanh^{-1}(x)\) 函数。t

位操作#

left_shift

逐个元素的左移。

right_shift

元素右移。

逻辑函数#

isnan

逐个元素的 NaN 检查。

isinf

逐个元素的 infinity 检查。

logical_and

Element-wise 逻辑与(AND)。

logical_not

Element-wise 逻辑非(NOT)。

logical_or

逐元素逻辑或操作。

logical_xor

逐元素逻辑异或。

greater

逐元素的大于比较。

greater_equal

逐元素的大于或等于比较。

less

逐元素的小于比较。

less_equal

逐元素的小于或等于比较。

equal

逐元素相等比较。

not_equal

逐元素的不等式比较。

统计函数#

sum

计算给定维度 (或坐标轴) 元素的求和结果。

cumsum

计算给定轴上的向量元素的累积和。

prod

计算张量元素在给定轴上的乘积。

mean

计算tensor在给定轴(或所有轴)上的平均值。

min

计算tensor在给定轴(或多个轴)上的最小值。

max

计算tensor在给定轴(或多个轴)上的最大值。

var

计算给定维度 (或坐标轴) 元素的方差

std

在给定的维度 (或坐标轴) 上计算元素的标准差。

norm

计算张量在给定轴上的范数。

normalize

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

参见

想要返回 min, max 的索引而不是元素值,请参考 Searching 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 (默认情况下)个最小元素。

度量函数#

topk_accuracy

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

神经网络运算#

卷积运算#

conv1d

一维卷积运算。

conv2d

二维卷积运算。

conv3d

三维卷积运算。

local_conv2d

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

conv_transpose2d

二维转置卷积运算。

conv_transpose3d

3D 转置卷积算子。

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

Element-wise max(x, 0).

relu6

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

prelu

逐元素进行 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

对输入应用 layer normalization。

local_response_norm

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

线性函数#

linear

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

multi-head attention functions#

Dropout 函数#

dropout

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

Sparse 函数#

one_hot

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

indexing_one_hot

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

embedding

应用查找表进行 embedding。

损失函数#

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

二维图像的批量仿射变换。

warp_perspective

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

roi_pooling

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

roi_align

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

nms

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

correlation

对输入进行相关性计算。

nvof

NVIDIA光流SDK的实现

量化函数#

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 函数#

extern_opr_subgraph

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

tensorrt_runtime_opr

cambricon_runtime_opr

将一个序列化后的寒武纪模型加载成 MegEngine 中的一个运行时算子。

atlas_runtime_opr

加载 Atlas 模型。

magicmind_runtime_opr

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