This matches the implementation in Denoising Diffusion Probabilistic Models: From Fairseq. Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of "Attention Is All You Need".
(timesteps, embedding_dim)
| 21 | |
| 22 | |
| 23 | def get_timestep_embedding(timesteps, embedding_dim): |
| 24 | """ |
| 25 | This matches the implementation in Denoising Diffusion Probabilistic Models: |
| 26 | From Fairseq. |
| 27 | Build sinusoidal embeddings. |
| 28 | This matches the implementation in tensor2tensor, but differs slightly |
| 29 | from the description in Section 3.5 of "Attention Is All You Need". |
| 30 | """ |
| 31 | assert len(timesteps.shape) == 1 |
| 32 | |
| 33 | half_dim = embedding_dim // 2 |
| 34 | emb = math.log(10000) / (half_dim - 1) |
| 35 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) |
| 36 | emb = emb.to(device=timesteps.device) |
| 37 | emb = timesteps.float()[:, None] * emb[None, :] |
| 38 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) |
| 39 | if embedding_dim % 2 == 1: # zero pad |
| 40 | emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) |
| 41 | return emb |
| 42 | |
| 43 | |
| 44 | def nonlinearity(x): |