You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
24 lines
868 B
24 lines
868 B
2 years ago
|
import math
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def timestep_embedding(timesteps, dim, max_period=10000):
|
||
|
|
"""
|
||
|
|
Create sinusoidal timestep embeddings.
|
||
|
|
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
||
|
|
These may be fractional.
|
||
|
|
:param dim: the dimension of the output.
|
||
|
|
:param max_period: controls the minimum frequency of the embeddings.
|
||
|
|
:return: an [N x dim] Tensor of positional embeddings.
|
||
|
|
"""
|
||
|
|
half = dim // 2
|
||
|
|
freqs = torch.exp(
|
||
|
|
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
||
|
|
).to(device=timesteps.device)
|
||
|
|
args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
|
||
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||
|
|
if dim % 2:
|
||
|
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||
|
|
return embedding
|