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