def setup_class(self):
self.dataset = build_dataset_from_name("cora")
self.data = Data.from_pyg_data(self.dataset[0])
self.num_nodes = self.data.num_nodes
self.num_edges = self.data.num_edges
self.num_features = self.data.num_features
print("Call Setup")
def fit(self, model, dataset):
self.data = Data.from_pyg_data(dataset[0])
self.train_loader = NeighborSampler(data=self.data,
mask=self.data.train_mask,
sizes=self.sample_size,
batch_size=self.batch_size,
num_workers=self.num_workers,
shuffle=True)
self.test_loader = NeighborSampler(data=self.data,
mask=None,
sizes=[-1],
batch_size=self.batch_size,
shuffle=False)
self.model = model.to(self.device)
self.model.set_data_device(self.device)
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.lr,
weight_decay=self.weight_decay)
best_model = self.train()
self.model = best_model
acc, loss = self._test_step()
return dict(Acc=acc["test"], ValAcc=acc["val"])
