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)
| 48 | return normalized_feat * style_std.expand(size) + style_mean.expand(size) |
| 49 | |
| 50 | def get_timestep_embedding(timesteps, embedding_dim): |
| 51 | """ |
| 52 | This matches the implementation in Denoising Diffusion Probabilistic Models: |
| 53 | From Fairseq. |
| 54 | Build sinusoidal embeddings. |
| 55 | This matches the implementation in tensor2tensor, but differs slightly |
| 56 | from the description in Section 3.5 of "Attention Is All You Need". |
| 57 | """ |
| 58 | assert len(timesteps.shape) == 1 |
| 59 | |
| 60 | half_dim = embedding_dim // 2 |
| 61 | emb = math.log(10000) / (half_dim - 1) |
| 62 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) |
| 63 | emb = emb.to(device=timesteps.device) |
| 64 | emb = timesteps.float()[:, None] * emb[None, :] |
| 65 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) |
| 66 | if embedding_dim % 2 == 1: # zero pad |
| 67 | emb = torch.nn.functional.pad(emb, (0,1,0,0)) |
| 68 | return emb |
| 69 | |
| 70 | |
| 71 | def nonlinearity(x): |