| 7 | |
| 8 | # This implementation is borrowed from IDR: https://github.com/lioryariv/idr |
| 9 | class SDFNetwork(nn.Module): |
| 10 | def __init__(self, |
| 11 | d_in, |
| 12 | d_out, |
| 13 | d_hidden, |
| 14 | n_layers, |
| 15 | skip_in=(4,), |
| 16 | multires=0, |
| 17 | bias=0.5, |
| 18 | scale=1, |
| 19 | geometric_init=True, |
| 20 | weight_norm=True, |
| 21 | inside_outside=False): |
| 22 | super(SDFNetwork, self).__init__() |
| 23 | |
| 24 | dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out] |
| 25 | |
| 26 | self.embed_fn_fine = None |
| 27 | |
| 28 | if multires > 0: |
| 29 | embed_fn, input_ch = get_embedder(multires, input_dims=d_in) |
| 30 | self.embed_fn_fine = embed_fn |
| 31 | dims[0] = input_ch |
| 32 | |
| 33 | self.num_layers = len(dims) |
| 34 | self.skip_in = skip_in |
| 35 | self.scale = scale |
| 36 | |
| 37 | for l in range(0, self.num_layers - 1): |
| 38 | if l + 1 in self.skip_in: |
| 39 | out_dim = dims[l + 1] - dims[0] |
| 40 | else: |
| 41 | out_dim = dims[l + 1] |
| 42 | |
| 43 | lin = nn.Linear(dims[l], out_dim) |
| 44 | |
| 45 | if geometric_init: |
| 46 | if l == self.num_layers - 2: |
| 47 | if not inside_outside: |
| 48 | torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001) |
| 49 | torch.nn.init.constant_(lin.bias, -bias) |
| 50 | else: |
| 51 | torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001) |
| 52 | torch.nn.init.constant_(lin.bias, bias) |
| 53 | elif multires > 0 and l == 0: |
| 54 | torch.nn.init.constant_(lin.bias, 0.0) |
| 55 | torch.nn.init.constant_(lin.weight[:, 3:], 0.0) |
| 56 | torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim)) |
| 57 | elif multires > 0 and l in self.skip_in: |
| 58 | torch.nn.init.constant_(lin.bias, 0.0) |
| 59 | torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim)) |
| 60 | torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0) |
| 61 | else: |
| 62 | torch.nn.init.constant_(lin.bias, 0.0) |
| 63 | torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim)) |
| 64 | |
| 65 | if weight_norm: |
| 66 | lin = nn.utils.weight_norm(lin) |