Code adapted from https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py
| 39 | |
| 40 | |
| 41 | class SinusoidalEmbedding(nn.Module): |
| 42 | """ Code adapted from https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py """ |
| 43 | |
| 44 | def __init__(self, embedding_dim, max_positions=10000, scale=1.0): |
| 45 | super().__init__() |
| 46 | self.embedding_dim = embedding_dim |
| 47 | self.max_positions = max_positions |
| 48 | self.scale = scale |
| 49 | |
| 50 | def forward(self, timesteps): |
| 51 | timesteps *= self.scale |
| 52 | assert timesteps.ndim == 1 |
| 53 | half_dim = self.embedding_dim // 2 |
| 54 | emb = math.log(self.max_positions) / (half_dim - 1) |
| 55 | emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb) |
| 56 | emb = timesteps.float()[:, None] * emb[None, :] |
| 57 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) |
| 58 | if self.embedding_dim % 2 == 1: # zero pad |
| 59 | emb = F.pad(emb, (0, 1), mode='constant') |
| 60 | assert emb.shape == (timesteps.shape[0], self.embedding_dim) |
| 61 | return emb |
| 62 | |
| 63 | |
| 64 | class GaussianFourierEmbedding(nn.Module): |
no outgoing calls
no test coverage detected