| 4 | |
| 5 | |
| 6 | class Complete_HCNN(pl.LightningModule): |
| 7 | def __init__(self, lighten, homological_structures): |
| 8 | super().__init__() |
| 9 | self.name = "hcnn" |
| 10 | if lighten: |
| 11 | self.name += "-lighten" |
| 12 | |
| 13 | self.homological_structures = homological_structures |
| 14 | self.tetrahedra = self.homological_structures['tetrahedra'] |
| 15 | self.triangles = self.homological_structures['triangles'] |
| 16 | self.edges = self.homological_structures['edges'] |
| 17 | |
| 18 | # ------------ # |
| 19 | |
| 20 | self.conv1_tetrahedra = nn.Sequential( |
| 21 | nn.Conv2d( |
| 22 | in_channels=1, out_channels=32, kernel_size=(1, 2), stride=(1, 2) |
| 23 | ), |
| 24 | nn.ReLU(), |
| 25 | ) |
| 26 | |
| 27 | self.conv1_triangles = nn.Sequential( |
| 28 | nn.Conv2d( |
| 29 | in_channels=1, out_channels=32, kernel_size=(1, 2), stride=(1, 2) |
| 30 | ), |
| 31 | nn.ReLU(), |
| 32 | ) |
| 33 | |
| 34 | self.conv1_edges = nn.Sequential( |
| 35 | nn.Conv2d( |
| 36 | in_channels=1, out_channels=32, kernel_size=(1, 2), stride=(1, 2) |
| 37 | ), |
| 38 | nn.ReLU(), |
| 39 | ) |
| 40 | |
| 41 | # ------------ # |
| 42 | |
| 43 | self.conv2_tetrahedra = nn.Sequential( |
| 44 | nn.Conv2d( |
| 45 | in_channels=32, out_channels=32, kernel_size=(1, 4), stride=(1, 4) |
| 46 | ), |
| 47 | nn.ReLU(), |
| 48 | nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(4, 1)), |
| 49 | nn.ReLU(), |
| 50 | nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(4, 1)), |
| 51 | nn.ReLU(), |
| 52 | ) |
| 53 | |
| 54 | self.conv2_triangles = nn.Sequential( |
| 55 | nn.Conv2d( |
| 56 | in_channels=32, out_channels=32, kernel_size=(1, 3), stride=(1, 3) |
| 57 | ), |
| 58 | nn.ReLU(), |
| 59 | nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(4, 1)), |
| 60 | nn.ReLU(), |
| 61 | nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(4, 1)), |
| 62 | nn.ReLU(), |
| 63 | ) |