(self, x)
| 69 | self._B = torch.stack(B_sort) # for sape |
| 70 | |
| 71 | def forward(self, x): |
| 72 | # assert x.dim() == 4, 'Expected 4D input (got {}D input)'.format(x.dim()) |
| 73 | |
| 74 | batches, channels = x.shape |
| 75 | |
| 76 | assert channels == self._num_input_channels, \ |
| 77 | "Expected input to have {} channels (got {} channels)".format(self._num_input_channels, channels) |
| 78 | |
| 79 | # Make shape compatible for matmul with _B. |
| 80 | # From [B, C, W, H] to [(B*W*H), C]. |
| 81 | # x = x.permute(0, 2, 3, 1).reshape(batches * width * height, channels) |
| 82 | |
| 83 | res = x @ self._B.to(x.device) |
| 84 | |
| 85 | # From [(B*W*H), C] to [B, W, H, C] |
| 86 | # x = x.view(batches, width, height, self._mapping_size) |
| 87 | # From [B, W, H, C] to [B, C, W, H] |
| 88 | # x = x.permute(0, 3, 1, 2) |
| 89 | |
| 90 | res = 2 * np.pi * res |
| 91 | return torch.cat([x,torch.sin(res), torch.cos(res)], dim=1) |
nothing calls this directly
no outgoing calls
no test coverage detected