| 20 | |
| 21 | class BlockSiren(torch.nn.Module): |
| 22 | def __init__(self, in_channels, out_channels, use_bias=True, is_first_layer=False, scale_init=1.0): |
| 23 | super(BlockSiren, self).__init__() |
| 24 | self.use_bias = use_bias |
| 25 | # self.activ = activ |
| 26 | self.is_first_layer = is_first_layer |
| 27 | self.scale_init = scale_init |
| 28 | # self.freq_scaling = scale_init |
| 29 | |
| 30 | self.conv = torch.nn.Linear(in_channels, out_channels, bias=self.use_bias).cuda() |
| 31 | |
| 32 | with torch.no_grad(): |
| 33 | |
| 34 | #following the official implementation from https://github.com/vsitzmann/siren/blob/master/explore_siren.ipynb |
| 35 | if self.is_first_layer: |
| 36 | self.conv.weight.uniform_(-1 / in_channels, |
| 37 | 1 / in_channels) |
| 38 | else: |
| 39 | self.conv.weight.uniform_(-np.sqrt(6 / in_channels) / self.scale_init, |
| 40 | np.sqrt(6 / in_channels) / self.scale_init) |
| 41 | |
| 42 | self.conv.bias.data.zero_() |
| 43 | |
| 44 | def forward(self, x): |
| 45 | x = self.conv(x) |