(weight, scale=1.0)
| 31 | return module |
| 32 | |
| 33 | def init_lecun_normal_param(weight, scale=1.0): |
| 34 | def truncated_normal(uniform, mu=0.0, sigma=1.0, a=-2, b=2): |
| 35 | normal = torch.distributions.normal.Normal(0, 1) |
| 36 | |
| 37 | alpha = (a - mu) / sigma |
| 38 | beta = (b - mu) / sigma |
| 39 | |
| 40 | alpha_normal_cdf = normal.cdf(torch.tensor(alpha)) |
| 41 | p = alpha_normal_cdf + (normal.cdf(torch.tensor(beta)) - alpha_normal_cdf) * uniform |
| 42 | |
| 43 | v = torch.clamp(2 * p - 1, -1 + 1e-8, 1 - 1e-8) |
| 44 | x = mu + sigma * np.sqrt(2) * torch.erfinv(v) |
| 45 | x = torch.clamp(x, a, b) |
| 46 | |
| 47 | return x |
| 48 | |
| 49 | def sample_truncated_normal(shape, scale=1.0): |
| 50 | stddev = np.sqrt(scale/shape[-1])/.87962566103423978 # shape[-1] = fan_in |
| 51 | return stddev * truncated_normal(torch.rand(shape)) |
| 52 | |
| 53 | weight = torch.nn.Parameter( (sample_truncated_normal(weight.shape)) ) |
| 54 | return weight |
| 55 | |
| 56 | # for gradient checkpointing |
| 57 | def create_custom_forward(module, **kwargs): |
no test coverage detected