13fromtypingimportDict1415importnumpyasnp16importtorch17fromtorchimportnn1819fromlabmlimportlab,monit,tracker,experiment20fromlabml.configsimportBaseConfigs,option,calculate21fromlabml.utilsimportdownload22fromlabml\_helpers.deviceimportDeviceConfigs23fromlabml\_helpers.moduleimportModule24fromlabml\_nn.graphs.gatimportGraphAttentionLayer25fromlabml\_nn.optimizers.configsimportOptimizerConfigs

28classCoraDataset:

43labels:torch.Tensor

45classes:Dict[str,int]

47features:torch.Tensor

50adj\_mat:torch.Tensor

52@staticmethod53def\_download():

57ifnot(lab.get\_data\_path()/'cora').exists():58download.download\_file('https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz',59lab.get\_data\_path()/'cora.tgz')60download.extract\_tar(lab.get\_data\_path()/'cora.tgz',lab.get\_data\_path())

62def\_\_init\_\_(self,include\_edges:bool=True):

69self.include\_edges=include\_edges

72self.\_download()

75withmonit.section('Read content file'):76content=np.genfromtxt(str(lab.get\_data\_path()/'cora/cora.content'),dtype=np.dtype(str))

78withmonit.section('Read citations file'):79citations=np.genfromtxt(str(lab.get\_data\_path()/'cora/cora.cites'),dtype=np.int32)

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

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

87self.classes={s:ifori,sinenumerate(set(content[:,-1]))}

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

92paper\_ids=np.array(content[:,0],dtype=np.int32)

94ids\_to\_idx={id\_:ifori,id\_inenumerate(paper\_ids)}

97self.adj\_mat=torch.eye(len(self.labels),dtype=torch.bool)

100ifself.include\_edges:101foreincitations:

103e1,e2=ids\_to\_idx[e[0]],ids\_to\_idx[e[1]]

107self.adj\_mat[e1][e2]=True108self.adj\_mat[e2][e1]=True

111classGAT(Module):

118def\_\_init\_\_(self,in\_features:int,n\_hidden:int,n\_classes:int,n\_heads:int,dropout:float):

126super().\_\_init\_\_()

129self.layer1=GraphAttentionLayer(in\_features,n\_hidden,n\_heads,is\_concat=True,dropout=dropout)

131self.activation=nn.ELU()

133self.output=GraphAttentionLayer(n\_hidden,n\_classes,1,is\_concat=False,dropout=dropout)

135self.dropout=nn.Dropout(dropout)

137defforward(self,x:torch.Tensor,adj\_mat:torch.Tensor):

144x=self.dropout(x)

146x=self.layer1(x,adj\_mat)

148x=self.activation(x)

150x=self.dropout(x)

152returnself.output(x,adj\_mat)

155defaccuracy(output:torch.Tensor,labels:torch.Tensor):

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

162classConfigs(BaseConfigs):

168model:GAT

170training\_samples:int=500

172in\_features:int

174n\_hidden:int=64

176n\_heads:int=8

178n\_classes:int

180dropout:float=0.6

182include\_edges:bool=True

184dataset:CoraDataset

186epochs:int=1\_000

188loss\_func=nn.CrossEntropyLoss()

193device:torch.device=DeviceConfigs()

195optimizer:torch.optim.Adam

197defrun(self):

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

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

211edges\_adj=self.dataset.adj\_mat.to(self.device)

213edges\_adj=edges\_adj.unsqueeze(-1)

216idx\_rand=torch.randperm(len(labels))

218idx\_train=idx\_rand[:self.training\_samples]

220idx\_valid=idx\_rand[self.training\_samples:]

223forepochinmonit.loop(self.epochs):

225self.model.train()

227self.optimizer.zero\_grad()

229output=self.model(features,edges\_adj)

231loss=self.loss\_func(output[idx\_train],labels[idx\_train])

233loss.backward()

235self.optimizer.step()

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

239tracker.add('accuracy.train',accuracy(output[idx\_train],labels[idx\_train]))

242self.model.eval()

245withtorch.no\_grad():

247output=self.model(features,edges\_adj)

249loss=self.loss\_func(output[idx\_valid],labels[idx\_valid])

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

253tracker.add('accuracy.valid',accuracy(output[idx\_valid],labels[idx\_valid]))

256tracker.save()

259@option(Configs.dataset)260defcora\_dataset(c:Configs):

264returnCoraDataset(c.include\_edges)

268calculate(Configs.n\_classes,lambdac:len(c.dataset.classes))

270calculate(Configs.in\_features,lambdac:c.dataset.features.shape[1])

273@option(Configs.model)274defgat\_model(c:Configs):

278returnGAT(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):

286opt\_conf=OptimizerConfigs()287opt\_conf.parameters=c.model.parameters()288returnopt\_conf

291defmain():

293conf=Configs()

295experiment.create(name='gat')

297experiment.configs(conf,{

299'optimizer.optimizer':'Adam',300'optimizer.learning\_rate':5e-3,301'optimizer.weight\_decay':5e-4,302})

305withexperiment.start():

307conf.run()

311if\_\_name\_\_=='\_\_main\_\_':312main()