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)
| 10 | |
| 11 | |
| 12 | def get_timestep_embedding(timesteps, embedding_dim): |
| 13 | """ |
| 14 | This matches the implementation in Denoising Diffusion Probabilistic Models: |
| 15 | From Fairseq. |
| 16 | Build sinusoidal embeddings. |
| 17 | This matches the implementation in tensor2tensor, but differs slightly |
| 18 | from the description in Section 3.5 of "Attention Is All You Need". |
| 19 | """ |
| 20 | assert len(timesteps.shape) == 1 |
| 21 | |
| 22 | half_dim = embedding_dim // 2 |
| 23 | emb = math.log(10000) / (half_dim - 1) |
| 24 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) |
| 25 | emb = emb.to(device=timesteps.device) |
| 26 | emb = timesteps.float()[:, None] * emb[None, :] |
| 27 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) |
| 28 | if embedding_dim % 2 == 1: # zero pad |
| 29 | emb = torch.nn.functional.pad(emb, (0,1,0,0)) |
| 30 | return emb |
| 31 | |
| 32 | |
| 33 | def nonlinearity(x): |