| 7 | |
| 8 | |
| 9 | class DecoderBase(ABC, nn.Module): |
| 10 | def __init__(self, conf): |
| 11 | super().__init__() |
| 12 | self.conf = conf |
| 13 | |
| 14 | hadarard_op = lambda lst: reduce(lambda x, y: x * y, lst) |
| 15 | summation_op = lambda lst: sum(lst) |
| 16 | self.op = hadarard_op if conf.model.hadamard else summation_op |
| 17 | |
| 18 | self.conv_block = HexplaneConvBlock(conf.model.latent_channels, conf.model.query_channels) |
| 19 | self.hex_upsample_block = None # subclass |
| 20 | |
| 21 | def forward(self, hexplane): |
| 22 | hexplane = self.conv_block(hexplane) # increase number of channels |
| 23 | hexplane = [block(plane) for block, plane in zip(self.hex_upsample_block, hexplane)] |
| 24 | coords = [[2, 5], [2, 4], [2, 3], [4, 5], [3, 5], [3, 4]] |
| 25 | hexplane = [plane.unsqueeze(dim1).unsqueeze(dim2).permute(0, 2, 3, 4, 5, 1) for (dim1, dim2), plane in |
| 26 | zip(coords, hexplane)] # unsqueeze |
| 27 | voxel = self.op(hexplane) # broadcast and combine |
| 28 | out = self.forward_voxel(voxel) |
| 29 | return out |
| 30 | |
| 31 | @abstractmethod |
| 32 | def forward_voxel(self, voxel): |
| 33 | pass |
| 34 | |
| 35 | |
| 36 | class ConvDecoder(DecoderBase): |
nothing calls this directly
no outgoing calls
no test coverage detected