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
(timesteps, dim, max_period=10000)
| 40 | |
| 41 | |
| 42 | def timestep_embedding(timesteps, dim, max_period=10000): |
| 43 | """ |
| 44 | Create sinusoidal timestep embeddings. |
| 45 | :param timesteps: a 1-D Tensor of N indices, one per batch element. |
| 46 | These may be fractional. |
| 47 | :param dim: the dimension of the output. |
| 48 | :param max_period: controls the minimum frequency of the embeddings. |
| 49 | :return: an [N x dim] Tensor of positional embeddings. |
| 50 | """ |
| 51 | half = dim // 2 |
| 52 | idx = torch.arange(start=0, end=half, dtype=torch.float32) |
| 53 | freqs = torch.exp(-math.log(max_period) * idx / |
| 54 | half).to(device=timesteps.device) |
| 55 | args = timesteps[:, None].float() * freqs[None] |
| 56 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| 57 | if dim % 2: |
| 58 | embedding = torch.cat( |
| 59 | [embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| 60 | return embedding |