NOTE the returned t is in range [0, 1], before sending into transformer, you need to scale it by num_train_steps = 1000
(
batch_size: int,
logit_mean: float = 0.0,
logid_std: float = 1.0,
device=None,
)
| 47 | |
| 48 | @staticmethod |
| 49 | def logitnorm_sample_t( |
| 50 | batch_size: int, |
| 51 | logit_mean: float = 0.0, |
| 52 | logid_std: float = 1.0, |
| 53 | device=None, |
| 54 | ): |
| 55 | """ |
| 56 | NOTE the returned t is in range [0, 1], before sending into transformer, you need to scale it by num_train_steps = 1000 |
| 57 | """ |
| 58 | t = torch.normal( |
| 59 | mean=logit_mean, |
| 60 | std=logid_std, |
| 61 | size=(batch_size, ), |
| 62 | device=device or 'cpu') |
| 63 | t = F.sigmoid(t) |
| 64 | return t |
| 65 | |
| 66 | def sigma_shift( |
| 67 | self, |