| 106 | |
| 107 | |
| 108 | class QueryNetwork(nn.Module): |
| 109 | def __init__(self): |
| 110 | super(QueryNetwork, self).__init__() |
| 111 | # TODO: Add proper non-linearity. |
| 112 | self.model = nn.Sequential( |
| 113 | nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 288), nn.ReLU(), nn.Linear(288, 288) |
| 114 | ) |
| 115 | |
| 116 | def forward(self, query): |
| 117 | out = self.model(query) |
| 118 | return out.reshape(-1, 4, 72) |
| 119 | |
| 120 | |
| 121 | class SpatialBasis: |