Shortcuts

hfai.nn.functional

set_replace_torch

把所有 hfai 优化过的 torch.nn.functional 和 torch 中的函数转换为 hfai.nn.functional 的对应函数.

dropout

hardtanh

hardtanh 函数, 参考 Hardtanh

hardtanh_

原地操作的 hardtanh 函数, 参考 Hardtanh

log_softmax

softmax

softmin

softmin 函数, 参考 Softmin

softplus

sync

allgather 输入的 tensor 并沿着指定的维度拼接在一起,支持 autograd,backward 的时候梯度会传回去

relu

relu_

原地操作的 relu 函数, 参考 ReLU

relu6

relu6 函数, 参考 ReLU6

threshold

threshold_

原地操作的 threshold 函数, 参考 Threshold

leaky_relu

leaky_relu_

原地操作的 leaky_relu 函数

rrelu

rrelu_

原地操作的 rrelu 函数

hardsigmoid

hardshrink

softshrink

abs

abs_

minimum

maximum

min

max

clip

clip 函数, 参考 hfai.nn.functional.clamp()

clip_

原地操作的 clip 函数

clamp

clamp_

原地操作的 clamp 函数

clamp_max

压位 clamp_max 算子, 训练时的中间结果用 1bit 储存 [x <= max], 以节省训练时的内存

clamp_max_

原地操作的 clamp_max 函数

clamp_min

压位 clamp_min 算子, 训练时的中间结果用 1bit 储存 [x >= min], 以节省训练时的内存

clamp_min_

原地操作的 clamp_min 函数

where

masked_fill

masked_fill_

masked_select

masked_scatter

masked_scatter_

class hfai.nn.functional.set_replace_torch(mode=True)[source]

把所有 hfai 优化过的 torch.nn.functional 和 torch 中的函数转换为 hfai.nn.functional 的对应函数. 如: torch.nn.functional.relu -> hfai.nn.functional.relu, torch.max -> hfai.nn.funtional.max

Parameters

mode (bool, optional) – 是否开启替换. 默认: True

Note

调用 hfai.nn.functional.set_replace_torch() 后, 无论是 (1)用户显示调用 还是 (2)PyTorch内部调用, 一切对 torch.nn.functional.xxx 的调用, 都会执行 hfai.nn.functional.xxx .

例如 torch.nn.CrossEntropyLoss 中的 log_softmax 会自动执行 hfai.nn.functional.log_softmax

Examples:

hfai.nn.functional.set_replace_torch()
y = torch.nn.functional.softmax(x) # softmax 执行 hfai 的实现
hfai.nn.functional.set_replace_torch(False)

with hfai.nn.functional.set_replace_torch():
    loss = torch.nn.functional.cross_entropy(input, target) # 内部的 log_softmax 执行 hfai 的实现
hfai.nn.functional.dropout(*args, **kwargs)
hfai.nn.functional.hardtanh(input, min_val=- 1.0, max_val=1.0, inplace=False)[source]

hardtanh 函数, 参考 Hardtanh

hfai.nn.functional.hardtanh_(input, min_val=- 1.0, max_val=1.0)[source]

原地操作的 hardtanh 函数, 参考 Hardtanh

hfai.nn.functional.log_softmax(*args, **kwargs)
hfai.nn.functional.softmax(*args, **kwargs)
hfai.nn.functional.softmin(input, dim=None, _stacklevel=3, dtype=None)[source]

softmin 函数, 参考 Softmin

hfai.nn.functional.softplus(*args, **kwargs)
hfai.nn.functional.sync(x, dist_group=False, dim=0, equal_size=False, tag=None, enable_timer=True, log_every_steps=1, timeout=60)[source]

allgather 输入的 tensor 并沿着指定的维度拼接在一起,支持 autograd,backward 的时候梯度会传回去

F.sync.get_metrics 会返回一个字典,格式如下:

{
    "tag1": {"iters": 100, "fwd": 25, "bwd": 40, "size": 16},
    "tag2": {"iters": 100, "fwd": 25, "bwd": 40, "size": 16},
}

iters 代表该 tag 调用的次数,fwd / bwd 代表每次 forward / backward 的平均耗时(ms),size 代表每次 forward 返回结果的平均大小(byte)

Parameters
  • x (Tensor) – 输入的 tensor

  • dist_group (ProcessGroup) – ProcessGroup 对象,如果是 False 则不会做 allgather

  • dim (int) – allgather 之后拼接的维度

  • equal_size (bool) – 是否每张卡上的 tensor 大小相同

  • tag (str) – 计时的标签,每个标签在一次 forward 中只能用一次; tag 为 None 时不计时

  • enable_timer (bool) – 是否计时

  • log_every_steps (int) – 每多少个 step 计时一次

  • timeout (int) – 本函数超时的秒数,超过这个时间会抛出异常;0 代表没有时间限制;默认是 60

Returns

拼接后的结果

Return type

out (Tensor)

Examples:

import torch.distributed as dist
import hfai.nn.functional as F

# init process group ...

rank = dist.get_rank()
x = torch.ones(1, requires_grad=True, device='cuda') * rank
out = F.sync(x, dist_group, dim=0, tag='tag1')
out.sum().backward()

# 打印耗时、通讯量等
F.sync.print_metrics()

# 获得 metrics
print(F.sync.get_metrics())

# 重置 metrics
F.sync.reset()
hfai.nn.functional.relu(*args, **kwargs)
hfai.nn.functional.relu_(input)[source]

原地操作的 relu 函数, 参考 ReLU

hfai.nn.functional.relu6(input, inplace=False)[source]

relu6 函数, 参考 ReLU6

hfai.nn.functional.threshold(*args, **kwargs)
hfai.nn.functional.threshold_(input, threshold, value)[source]

原地操作的 threshold 函数, 参考 Threshold

hfai.nn.functional.leaky_relu(*args, **kwargs)
hfai.nn.functional.leaky_relu_(input, negative_slope=0.01)[source]

原地操作的 leaky_relu 函数

hfai.nn.functional.rrelu(*args, **kwargs)
hfai.nn.functional.rrelu_(input, lower=0.125, upper=0.3333333333333333, training=False)[source]

原地操作的 rrelu 函数

hfai.nn.functional.hardsigmoid(*args, **kwargs)
hfai.nn.functional.hardshrink(*args, **kwargs)
hfai.nn.functional.softshrink(*args, **kwargs)
hfai.nn.functional.abs(*args, **kwargs)
hfai.nn.functional.abs_(*args, **kwargs)
hfai.nn.functional.minimum(*args, **kwargs)
hfai.nn.functional.maximum(*args, **kwargs)
hfai.nn.functional.min(*args, **kwargs)
hfai.nn.functional.max(*args, **kwargs)
hfai.nn.functional.clip(input, min=None, max=None)[source]

clip 函数, 参考 hfai.nn.functional.clamp()

hfai.nn.functional.clip_(input, min=None, max=None)[source]

原地操作的 clip 函数

hfai.nn.functional.clamp(*args, **kwargs)
hfai.nn.functional.clamp_(input, min=None, max=None)[source]

原地操作的 clamp 函数

hfai.nn.functional.clamp_max(input, max, inplace=False)[source]

压位 clamp_max 算子, 训练时的中间结果用 1bit 储存 [x <= max], 以节省训练时的内存

Parameters
  • input – 输入的 Tensor

  • max (float) – output 的最大值

  • inplace (bool, optional) – 如果是 True, 进行原地操作, 默认: False

import hfai.nn.functional as F

y = F.clamp_max(x, max=0.5)
# same as: y = x.clamp_max(max=0.5)
# same as: y = torch.min(x, 0.5 * torch.ones_like(x))
hfai.nn.functional.clamp_max_(input, max)[source]

原地操作的 clamp_max 函数

hfai.nn.functional.clamp_min(input, min, inplace=False)[source]

压位 clamp_min 算子, 训练时的中间结果用 1bit 储存 [x >= min], 以节省训练时的内存

Parameters
  • input – 输入的 Tensor

  • min (float) – output 的最小值

  • inplace (bool, optional) – 如果是 True, 进行原地操作, 默认: False

import hfai.nn.functional as F

y = F.clamp_min(x, min=-0.5)
# same as: y = x.clamp_min(min=-0.5)
# same as: y = torch.max(x, -0.5 * torch.ones_like(x))
hfai.nn.functional.clamp_min_(input, min)[source]

原地操作的 clamp_min 函数

hfai.nn.functional.where(*args, **kwargs)
hfai.nn.functional.masked_fill(*args, **kwargs)
hfai.nn.functional.masked_fill_(*args, **kwargs)
hfai.nn.functional.masked_select(*args, **kwargs)
hfai.nn.functional.masked_scatter(*args, **kwargs)
hfai.nn.functional.masked_scatter_(*args, **kwargs)