Create sinusoidal timestep embeddings. :param t: 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.
(t, dim, max_period=10000)
| 207 | |
| 208 | @staticmethod |
| 209 | def timestep_embedding(t, dim, max_period=10000): |
| 210 | """ |
| 211 | Create sinusoidal timestep embeddings. |
| 212 | :param t: a 1-D Tensor of N indices, one per batch element. |
| 213 | These may be fractional. |
| 214 | :param dim: the dimension of the output. |
| 215 | :param max_period: controls the minimum frequency of the embeddings. |
| 216 | :return: an (N, D) Tensor of positional embeddings. |
| 217 | """ |
| 218 | # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py |
| 219 | half = dim // 2 |
| 220 | freqs = torch.exp( |
| 221 | -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half |
| 222 | ).to(device=t.device) |
| 223 | args = t[:, None].float() * freqs[None] |
| 224 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| 225 | if dim % 2: |
| 226 | embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| 227 | return embedding |
| 228 | |
| 229 | def forward(self, t, dtype=torch.float32): |
| 230 | t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype) |