gistfile1.txt

import time

import torch
import torch.nn as nn

from apex.normalization import FusedLayerNorm


torch.backends.cudnn.benchmark = True

nb_iters = 10000

# Create data
x = torch.randn(512, 16, 1024, device='cuda')

# upstream layernorm
norm = nn.LayerNorm(x.size()[-1]).cuda()

# apex fusedlayernorm
fused_norm = FusedLayerNorm(x.size()[-1]).cuda()


def run(name, module, input_):
    # cudnn warmup
    for _ in range(50):
        _ = module(input_)

    torch.cuda.synchronize()
    t0 = time.time()

    for _ in range(nb_iters):
        _ = module(input_)

    torch.cuda.synchronize()
    t1 = time.time()

    print("{} layernorm {:.3f}".format(name, t1 - t0))

run('upstream', norm, x)
run('apex', fused_norm, x)

x = x.half()
run('upstream half', norm.half(), x)
run('apex half', fused_norm.half(), x)