Pytorch activation function speed

README.md

Performance for activation function

timeit applied to various Pytorch activation function on CPU and GPU, on a tensor of shape (20000,)

You must run this in a IPython session. Pytorch version is 2.0.0+cu117.

See run.py.

Results

Activation	GPU time	CPU time
Mish	7.78 µs ± 73.7 ns	107 µs ± 2.83 µs
ReLU	8.15 µs ± 61 ns	7.54 µs ± 107 ns
SiLU	7.93 µs ± 16.9 ns	16.4 µs ± 67.9 ns
Sigmoid	6.49 µs ± 51.1 ns	14.6 µs ± 49.1 ns
Tanh	6.47 µs ± 49.4 ns	13.6 µs ± 351 ns
ELU	8.54 µs ± 55.3 ns	15.5 µs ± 356 ns
CELU	8.59 µs ± 38.9 ns	15.3 µs ± 78.7 ns
SELU	7.85 µs ± 72.7 ns	14.5 µs ± 78.4 ns
GELU	7.83 µs ± 117 ns	28 µs ± 6.88 µs
LeakyReLU	8.53 µs ± 59.6 ns	8.6 µs ± 341 ns
LogSigmoid	7.21 µs ± 139 ns	27.9 µs ± 482 ns
Softplus	8.59 µs ± 52.7 ns	44.7 µs ± 596 ns
Softsign	13.7 µs ± 57.7 ns	14.4 µs ± 217 ns
Tanhshrink	9.53 µs ± 50.3 ns	19.5 µs ± 283 ns
Softmin	14 µs ± 262 ns	22 µs ± 354 ns
Softmax	12.4 µs ± 212 ns	19.2 µs ± 280 ns
LogSoftmax	10.2 µs ± 64.4 ns	18.8 µs ± 307 ns

run.py

import torch
from torch import nn

fns = [
    nn.CELU, nn.ELU, nn.GELU, nn.LeakyReLU, nn.LogSigmoid, nn.LogSoftmax, nn.Mish, nn.ReLU, nn.SELU,
    nn.Sigmoid, nn.SiLU, nn.Softmax, nn.Softmin, nn.Softplus, nn.Softsign, nn.Tanh, nn.Tanhshrink
]


for dd in ["cuda", "cpu"]:
  x = torch.arange(-10, 10, 0.001, device=dd)  # shape=(20000,)
  for ff in fns:
        print(dd, ff.__qualname__)
        %timeit -r7 -n10000 _ = ff()(x)