කෝරාදත්ත කට්ටලය මත ප්රස්තාරය අවධානය ජාලය (GAT) පුහුණු

13from typing import Dict
14
15import numpy as np
16import torch
17from torch import nn
18
19from labml import lab, monit, tracker, experiment
20from labml.configs import BaseConfigs, option, calculate
21from labml.utils import download
22from labml_helpers.device import DeviceConfigs
23from labml_helpers.module import Module
24from labml_nn.graphs.gat import GraphAttentionLayer
25from labml_nn.optimizers.configs import OptimizerConfigs

28class CoraDataset:

43    labels: torch.Tensor

45    classes: Dict[str, int]

47    features: torch.Tensor

50    adj_mat: torch.Tensor

52    @staticmethod
53    def _download():

57        if not (lab.get_data_path() / 'cora').exists():
58            download.download_file('https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz',
59                                   lab.get_data_path() / 'cora.tgz')
60            download.extract_tar(lab.get_data_path() / 'cora.tgz', lab.get_data_path())

62    def __init__(self, include_edges: bool = True):

69        self.include_edges = include_edges

72        self._download()

75        with monit.section('Read content file'):
76            content = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.content'), dtype=np.dtype(str))

78        with monit.section('Read citations file'):
79            citations = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.cites'), dtype=np.int32)

82        features = torch.tensor(np.array(content[:, 1:-1], dtype=np.float32))

84        self.features = features / features.sum(dim=1, keepdim=True)

87        self.classes = {s: i for i, s in enumerate(set(content[:, -1]))}

89        self.labels = torch.tensor([self.classes[i] for i in content[:, -1]], dtype=torch.long)

92        paper_ids = np.array(content[:, 0], dtype=np.int32)

94        ids_to_idx = {id_: i for i, id_ in enumerate(paper_ids)}

97        self.adj_mat = torch.eye(len(self.labels), dtype=torch.bool)

100        if self.include_edges:
101            for e in citations:

103                e1, e2 = ids_to_idx[e[0]], ids_to_idx[e[1]]

107                self.adj_mat[e1][e2] = True
108                self.adj_mat[e2][e1] = True

111class GAT(Module):

118    def __init__(self, in_features: int, n_hidden: int, n_classes: int, n_heads: int, dropout: float):

126        super().__init__()

129        self.layer1 = GraphAttentionLayer(in_features, n_hidden, n_heads, is_concat=True, dropout=dropout)

131        self.activation = nn.ELU()

133        self.output = GraphAttentionLayer(n_hidden, n_classes, 1, is_concat=False, dropout=dropout)

135        self.dropout = nn.Dropout(dropout)

137    def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):

144        x = self.dropout(x)

146        x = self.layer1(x, adj_mat)

148        x = self.activation(x)

150        x = self.dropout(x)

152        return self.output(x, adj_mat)

155def accuracy(output: torch.Tensor, labels: torch.Tensor):

159    return output.argmax(dim=-1).eq(labels).sum().item() / len(labels)

162class Configs(BaseConfigs):

168    model: GAT

170    training_samples: int = 500

172    in_features: int

174    n_hidden: int = 64

176    n_heads: int = 8

178    n_classes: int

180    dropout: float = 0.6

182    include_edges: bool = True

184    dataset: CoraDataset

186    epochs: int = 1_000

188    loss_func = nn.CrossEntropyLoss()

193    device: torch.device = DeviceConfigs()

195    optimizer: torch.optim.Adam

197    def run(self):

207        features = self.dataset.features.to(self.device)

209        labels = self.dataset.labels.to(self.device)

211        edges_adj = self.dataset.adj_mat.to(self.device)

213        edges_adj = edges_adj.unsqueeze(-1)

216        idx_rand = torch.randperm(len(labels))

218        idx_train = idx_rand[:self.training_samples]

220        idx_valid = idx_rand[self.training_samples:]

223        for epoch in monit.loop(self.epochs):

225            self.model.train()

227            self.optimizer.zero_grad()

229            output = self.model(features, edges_adj)

231            loss = self.loss_func(output[idx_train], labels[idx_train])

233            loss.backward()

235            self.optimizer.step()

237            tracker.add('loss.train', loss)

239            tracker.add('accuracy.train', accuracy(output[idx_train], labels[idx_train]))

242            self.model.eval()

245            with torch.no_grad():

247                output = self.model(features, edges_adj)

249                loss = self.loss_func(output[idx_valid], labels[idx_valid])

251                tracker.add('loss.valid', loss)

253                tracker.add('accuracy.valid', accuracy(output[idx_valid], labels[idx_valid]))

256            tracker.save()

259@option(Configs.dataset)
260def cora_dataset(c: Configs):

264    return CoraDataset(c.include_edges)

268calculate(Configs.n_classes, lambda c: len(c.dataset.classes))

270calculate(Configs.in_features, lambda c: c.dataset.features.shape[1])

273@option(Configs.model)
274def gat_model(c: Configs):

278    return GAT(c.in_features, c.n_hidden, c.n_classes, c.n_heads, c.dropout).to(c.device)

281@option(Configs.optimizer)
282def _optimizer(c: Configs):

286    opt_conf = OptimizerConfigs()
287    opt_conf.parameters = c.model.parameters()
288    return opt_conf

291def main():

293    conf = Configs()

295    experiment.create(name='gat')

297    experiment.configs(conf, {

299        'optimizer.optimizer': 'Adam',
300        'optimizer.learning_rate': 5e-3,
301        'optimizer.weight_decay': 5e-4,
302    })

305    with experiment.start():

307        conf.run()

311if __name__ == '__main__':
312    main()