Skip to content

Instantly share code, notes, and snippets.

@FindHao
Created July 26, 2022 17:38
Show Gist options
  • Save FindHao/abbaa7a8a0b74173d9eb2233684bb6a9 to your computer and use it in GitHub Desktop.
Save FindHao/abbaa7a8a0b74173d9eb2233684bb6a9 to your computer and use it in GitHub Desktop.
import torch
from torch import profiler
import numpy as np
def _len_and_dim_norm(vectors):
"""
length and attention head size dim normalization
"""
vectors = vectors * torch.rsqrt(
torch.tensor(64, device=vectors.device, dtype=vectors.dtype)
)
return vectors
def _len_and_dim_norm2(vectors):
"""
length and attention head size dim normalization
"""
# tmp = torch.tensor(64, device=vectors.device, dtype=vectors.dtype)
vectors = vectors / np.sqrt(64)
return vectors
def profile():
activity_groups = []
activity_groups.append(profiler.ProfilerActivity.CUDA)
activity_groups.append(profiler.ProfilerActivity.CPU)
profile_detailed = True
input_shape = [8, 12, 64, 64, 64]
input = torch.ones(input_shape, dtype=torch.float32, device='cuda')
with profiler.profile(
schedule=profiler.schedule(wait=0, warmup=0, active=1),
activities=activity_groups,
record_shapes=profile_detailed,
profile_memory=profile_detailed,
with_stack=profile_detailed,
with_flops=profile_detailed,
on_trace_ready=profiler.tensorboard_trace_handler('/tmp/logs/')
) as prof:
x = _len_and_dim_norm(input)
y = _len_and_dim_norm2(input)
return y
profile()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment