| 111 | |
| 112 | class MESH_NET(nn.Module): |
| 113 | def __init__(self): |
| 114 | super(MESH_NET,self).__init__() |
| 115 | self.feature_dim = 3 |
| 116 | self.conv1 = GCNConv(self.feature_dim, 32) |
| 117 | self.pool1 = TopKPooling(32, ratio=0.6) |
| 118 | self.conv2 = GCNConv(32, 32) #(32, 64) |
| 119 | self.pool2 = TopKPooling(32, ratio=0.6) #64, ratio=0.6) |
| 120 | self.conv3 = GCNConv(32, 32) #(64, 128) |
| 121 | self.pool3 = TopKPooling(32, ratio=0.6) #(128, ratio=0.6) |
| 122 | # self.item_embedding = torch.nn.Embedding(num_embeddings=df.item_id.max() +1, embedding_dim=self.feature_dim) |
| 123 | self.lin1 = torch.nn.Linear(64, 16) #(256, 128) |
| 124 | self.lin2 = torch.nn.Linear(16, 8) #(128, 64) |
| 125 | # self.lin3 = torch.nn.Linear(8, 1) #(64, 1) |
| 126 | self.bn1 = torch.nn.BatchNorm1d(16) #(128) |
| 127 | self.bn2 = torch.nn.BatchNorm1d(8) #(64) |
| 128 | self.act1 = torch.nn.ReLU() |
| 129 | # self.act2 = torch.nn.ReLU() |
| 130 | |
| 131 | def forward(self, data): |
| 132 | x, edge_index, batch = data.pos, data.edge_index, data.batch |