| 25 | |
| 26 | |
| 27 | class SAGE(nn.Module): |
| 28 | def __init__(self, in_size, hidden_size, out_size, num_layers, dropout): |
| 29 | super().__init__() |
| 30 | self.layers = nn.ModuleList() |
| 31 | # Three-layer GraphSAGE-mean. |
| 32 | self.layers.append(dglnn.SAGEConv(in_size, hidden_size, "mean")) |
| 33 | for _ in range(num_layers - 2): |
| 34 | self.layers.append(dglnn.SAGEConv(hidden_size, hidden_size, "mean")) |
| 35 | self.layers.append(dglnn.SAGEConv(hidden_size, out_size, "mean")) |
| 36 | self.dropout = nn.Dropout(dropout) |
| 37 | self.hidden_size = hidden_size |
| 38 | self.out_size = out_size |
| 39 | # Set the dtype for the layers manually. |
| 40 | self.set_layer_dtype(torch.float32) |
| 41 | |
| 42 | def set_layer_dtype(self, _dtype): |
| 43 | for layer in self.layers: |
| 44 | for param in layer.parameters(): |
| 45 | param.data = param.data.to(_dtype) |
| 46 | |
| 47 | def forward(self, blocks, x): |
| 48 | hidden_x = x |
| 49 | for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)): |
| 50 | hidden_x = layer(block, hidden_x) |
| 51 | is_last_layer = layer_idx == len(self.layers) - 1 |
| 52 | if not is_last_layer: |
| 53 | hidden_x = F.relu(hidden_x) |
| 54 | hidden_x = self.dropout(hidden_x) |
| 55 | return hidden_x |
| 56 | |
| 57 | def inference(self, graph, features, dataloader, storage_device): |
| 58 | """Conduct layer-wise inference to get all the node embeddings.""" |
| 59 | pin_memory = storage_device == "pinned" |
| 60 | buffer_device = torch.device("cpu" if pin_memory else storage_device) |
| 61 | |
| 62 | for layer_idx, layer in enumerate(self.layers): |
| 63 | is_last_layer = layer_idx == len(self.layers) - 1 |
| 64 | |
| 65 | y = torch.empty( |
| 66 | graph.total_num_nodes, |
| 67 | self.out_size if is_last_layer else self.hidden_size, |
| 68 | dtype=torch.float32, |
| 69 | device=buffer_device, |
| 70 | pin_memory=pin_memory, |
| 71 | ) |
| 72 | for data in tqdm(dataloader): |
| 73 | # len(blocks) = 1 |
| 74 | hidden_x = layer(data.blocks[0], data.node_features["feat"]) |
| 75 | if not is_last_layer: |
| 76 | hidden_x = F.relu(hidden_x) |
| 77 | hidden_x = self.dropout(hidden_x) |
| 78 | # By design, our output nodes are contiguous. |
| 79 | y[data.seeds[0] : data.seeds[-1] + 1] = hidden_x.to( |
| 80 | buffer_device |
| 81 | ) |
| 82 | if not is_last_layer: |
| 83 | features.update("node", None, "feat", y) |
| 84 | |