Could you also add a GPU implementation for it? We try to keep parity between CPU and CUDA as much as possible.

Also, what's the story around the JIT fusion for CPU? It's disabled by default in PyTorch, but the fuser is currently able to handle all the operations in here inside a single fusion group, and seems to bring some significant speedups for CPU (and the same also applies to CUDA)

import torch

# should we enable it by default?
torch._C._jit_override_can_fuse_on_cpu(True)

def gelu(x):
    sqrt_two = 1.4142135623730951
    return x * 0.5 * (1.0 + torch.erf(x / sqrt_two))

@torch.jit.script
def gelu2(x):
    sqrt_two = 1.4142135623730951
    return x * 0.5 * (1.0 + torch.erf(x / sqrt_two))

x = torch.rand(64, 128, 56, 56)

# compile gelu2
gelu2(x)

and gives

%timeit gelu(x)
> 92.6 ms ± 647 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit gelu2(x)
> 18.5 ms ± 518 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

For the current version, the MKL path in the same machine need 109ms while the non-MKL path need 112ms, currently they are quite similar.

This pull request has been merged in 93ae040.

Is this implemented in 1.2.0? I can find it in documentation (https://pytorch.org/docs/stable/nn.functional.html) but I can't import it or find it in my installed library.

@BramVanroy - is this it?

@BramVanroy - is this it?

Odd. My IDE (PyCharm) underlines gelu in red and says "Cannot find reference 'gelu' in functional.pyi", but when I run the code it seems to import just fine.

from torch.nn.functional import gelu

PyTorch's current implementation is

def gelu(x):
    return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))

rather than

def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(math.pi / 2) * (x + 0.044715 * x ** 3)))

Correct? I've seen both being mentioned, but I'm not sure which one is implemented in PyTorch. IIRC Google's BERT originally uses the former, and OpenAI's GPT the latter.

PyTorch implementation is the original definition of GELU which is x * P(X <= x) where X ~ N(0, 1). This one is mathematically equivalent to 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0))). The tanh is an approximation of GELU which may lead to better performance as mentioned in https://arxiv.org/pdf/1606.08415.pdf. However, in our test, the performance improvement depends on how tanh function is implemented. So we didn't use that by default.

@BIT-silence Thank you for the reply. That clarifies things a lot. A final question, is there a reason that gelu doesn't have a Module equivalent that lives in nn? (Like nn.ReLU.)

Actually the main reason is we don't have enough time when adding it. We will consider add the Module later.

Okay, thanks for the information. I wasn't sure whether in general there is a reason when an activation function has a Module equivalent, but it seems that it's mostly a time-constraint.

@BramVanroy generally we only keep functionals for layers which don't have learnable parameters. We used to add layers for all common functions, like nn.ReLU but that's legacy

@BramVanroy generally we only keep functionals for layers which don't have learnable parameters. We used to add layers for all common functions, like nn.ReLU but that's legacy

That makes sense. However, I do like that printing a Module, gives you a nice overview of its submodules. If the activation is a Module it is included (and clear to the user what it is). If it's a function, it's not included unfortunately - so you're left wondering whether (and where) an activation is taking place. Compare the two snippets below

# as Module
from torch import nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.dense = nn.Linear(512, 1)
        self.activation = nn.ReLU()

    def forward(self, inputs):
        out = self.dense(inputs)
        out = self.activation(out)
        return out


net = Net()
print(net)

Prints

Net(
  (dense): Linear(in_features=512, out_features=1, bias=True)
  (activation): ReLU()
)

But when using functions, you don't see which activations are used nor where.

from torch import nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.dense = nn.Linear(512, 1)
        self.activation = F.relu

    def forward(self, inputs):
        out = self.dense(inputs)
        out = self.activation(out)
        return out


net = Net()
print(net)

Prints

Net(
  (dense): Linear(in_features=512, out_features=1, bias=True)
)

The module which is torch.nn.GELU is implemented in #28944

Great! It'll ship with 1.4 then, I presume?