(self, num_drags, fourier_freqs=8, downsample_ratio=64)
| 22 | |
| 23 | class DragPositionNet(nn.Module): |
| 24 | def __init__(self, num_drags, fourier_freqs=8, downsample_ratio=64): |
| 25 | super().__init__() |
| 26 | self.num_drags = num_drags |
| 27 | |
| 28 | self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) |
| 29 | self.position_dim = fourier_freqs*2*2 # 2 for sin and cos, 2 for 2 dims (x1, y1) or (x2, y2) |
| 30 | |
| 31 | # -------------------------------------------------------------- # |
| 32 | self.linears_drag = nn.Sequential( |
| 33 | nn.Linear(self.position_dim, 128), |
| 34 | nn.SiLU(), |
| 35 | nn.Linear(128, 256), |
| 36 | nn.SiLU(), |
| 37 | nn.Linear(256, 512), |
| 38 | ) |
| 39 | |
| 40 | self.downsample_ratio = downsample_ratio |
| 41 | |
| 42 | |
| 43 | def forward(self, drags_start, drags_end): |
nothing calls this directly
no test coverage detected