| 5 | |
| 6 | |
| 7 | class LinModule(nn.Module): |
| 8 | def __init__(self, |
| 9 | d_in, |
| 10 | d_out, |
| 11 | dims, |
| 12 | multires=0, |
| 13 | act_fun=None, last_act_fun=None, weight_norm=False, weight_zero=False, weight_xavier=True): |
| 14 | super().__init__() |
| 15 | |
| 16 | dims = [d_in] + dims + [d_out] |
| 17 | self.num_layers = len(dims) |
| 18 | if act_fun is None: |
| 19 | self.act_fun = nn.Softplus(beta=100) |
| 20 | else: |
| 21 | self.act_fun = act_fun |
| 22 | self.last_act_fun = last_act_fun |
| 23 | |
| 24 | for l in range(0, self.num_layers - 1): |
| 25 | out_dim = dims[l + 1] |
| 26 | lin = nn.Linear(dims[l], out_dim) |
| 27 | |
| 28 | if multires > 0 and l == 0: |
| 29 | torch.nn.init.constant_(lin.bias, 0.0) |
| 30 | torch.nn.init.constant_(lin.weight[:, 3:], 0.0) |
| 31 | torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim)) |
| 32 | else: |
| 33 | torch.nn.init.constant_(lin.bias, 0.0) |
| 34 | torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim)) |
| 35 | if weight_zero: |
| 36 | torch.nn.init.normal_(lin.weight, 0.0, 0.0) |
| 37 | if weight_norm: |
| 38 | lin = nn.utils.weight_norm(lin) |
| 39 | if weight_xavier: |
| 40 | torch.nn.init.xavier_normal_(lin.weight) |
| 41 | torch.nn.init.constant_(lin.bias, 0.0) |
| 42 | |
| 43 | setattr(self, f"lin{l}", lin) |
| 44 | |
| 45 | def forward(self, inx): |
| 46 | x = inx |
| 47 | for l in range(self.num_layers - 1): |
| 48 | lin = getattr(self, f"lin{l}") |
| 49 | x = lin(x) |
| 50 | if l == self.num_layers - 2: |
| 51 | if self.last_act_fun is not None: |
| 52 | x = self.last_act_fun(x) |
| 53 | else: |
| 54 | x = self.act_fun(x) |
| 55 | return x |