def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.save_hyperparameters() self.model = DeepGraphInfomax( hidden_channels=kwargs["hidden_channels"], corruption=corruption, encoder=Encoder(kwargs["num_features"], kwargs["hidden_channels"]), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)))
def __init__(self, num_features, num_classes, hidden_size): super(Infomax, self).__init__() self.num_features = num_features self.num_classes = num_classes self.hidden_size = hidden_size self.model = DeepGraphInfomax( hidden_channels=hidden_size, encoder=Encoder(num_features, hidden_size), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption)
def test_deep_graph_infomax(): def corruption(z): return z + 1 model = DeepGraphInfomax(hidden_channels=16, encoder=lambda x: x, summary=lambda z, *args: z.mean(dim=0), corruption=lambda x: x + 1) assert model.__repr__() == 'DeepGraphInfomax(16)' x = torch.ones(20, 16) pos_z, neg_z, summary = model(x) assert pos_z.size() == (20, 16) and neg_z.size() == (20, 16) assert summary.size() == (16, ) loss = model.loss(pos_z, neg_z, summary) assert 0 <= loss.item() acc = model.test(torch.ones(20, 16), torch.randint(10, (20, )), torch.ones(20, 16), torch.randint(10, (20, ))) assert 0 <= acc and acc <= 1
class Infomax(BaseModel): @staticmethod def add_args(parser): """Add model-specific arguments to the parser.""" # fmt: off parser.add_argument("--num-features", type=int) parser.add_argument("--num-classes", type=int) parser.add_argument("--hidden-size", type=int, default=512) # fmt: on @classmethod def build_model_from_args(cls, args): return cls( args.num_features, args.num_classes, args.hidden_size, ) def __init__(self, num_features, num_classes, hidden_size): super(Infomax, self).__init__() self.num_features = num_features self.num_classes = num_classes self.hidden_size = hidden_size self.model = DeepGraphInfomax( hidden_channels=hidden_size, encoder=Encoder(num_features, hidden_size), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption, ) def forward(self, x, edge_index): return self.model(x, edge_index) def node_classification_loss(self, data): pos_z, neg_z, summary = self.forward(data.x, data.edge_index) loss = self.model.loss(pos_z, neg_z, summary) return loss def predict(self, data): z, _, _ = self.forward(data.x, data.edge_index) clf = LogisticRegression(solver="lbfgs", multi_class="auto", max_iter=150) clf.fit(z[data.train_mask].detach().cpu().numpy(), data.y[data.train_mask].detach().cpu().numpy()) logits = torch.Tensor(clf.predict_proba(z.detach().cpu().numpy())) if z.is_cuda: logits = logits.cuda() return logits
class DGILearner: def __init__(self, inp_dim, out_dim, device): self.encoder = DGIEncoderNet(inp_dim, out_dim) self.dgi = DeepGraphInfomax(out_dim, encoder=self.encoder, summary=self.readout, corruption=self.corrupt) self.dgi = self.dgi.to(device) self.optimizer = torch.optim.Adam(self.dgi.parameters()) def embed(self, data): pos_z, _, _ = self.dgi(data.x, data.edge_index, data.edge_attr, msk=None) return pos_z def readout(self, z, x, edge_index, edge_attr, msk=None): if msk is None: return torch.sigmoid(torch.mean(z, 0)) else: return torch.sigmoid(torch.sum(z[msk], 0) / torch.sum(msk)) def corrupt(self, x, edge_index, edge_attr, msk=None): shuffled_rows = torch.randperm(len(x)) shuffled_x = x[shuffled_rows, :] return shuffled_x, edge_index, edge_attr def evaluate_loss(self, data, mode): # use masking for loss evaluation pos_z_train, neg_z_train, summ_train = self.dgi(data.x, data.edge_index, data.edge_attr, msk=data.train_mask) pos_z_test, neg_z_test, summ_test = self.dgi(data.x, data.edge_index, data.edge_attr, msk=data.test_mask) if mode == 'train': return self.dgi.loss(pos_z_train, neg_z_train, summ_train) else: return self.dgi.loss(pos_z_test, neg_z_test, summ_test) def train(self, data): # training self.dgi.train() self.optimizer.zero_grad() loss = self.evaluate_loss(data, mode='train') loss.backward() self.optimizer.step() return loss.item() def test(self, data): # testing self.dgi.eval() return self.evaluate_loss(data, mode='test').item()
class InfoMaxModel(BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.save_hyperparameters() self.model = DeepGraphInfomax( hidden_channels=kwargs["hidden_channels"], corruption=corruption, encoder=Encoder(kwargs["num_features"], kwargs["hidden_channels"]), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0))) def forward(self, batch): pos_z, neg_z, summary = self.model(batch.x, batch.edge_index[0]) loss = self.model.loss(pos_z, neg_z, summary) return pos_z, loss
def forward(self, x, adjs): for i, (edge_index, _, size) in enumerate(adjs): x_target = x[:size[1]] # Target nodes are always placed first. x = self.convs[i]((x, x_target), edge_index) x = self.activations[i](x) return x def corruption(x, edge_index): return x[torch.randperm(x.size(0))], edge_index device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = DeepGraphInfomax( hidden_channels=512, encoder=Encoder(dataset.num_features, 512), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption).to(device) model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.0001) x, y = data.x.to(device), data.y.to(device) def train(epoch): model.train() total_loss = total_examples = 0 for batch_size, n_id, adjs in tqdm(train_loader, desc=f'Epoch {epoch:02d}'): # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
def infomax(fp, PARAMS, feature): """[generate DGI embedding] Args: fp ([string]): [file path of the root of the data] PARAMS ([dict]): [the parameters of the node2vec model, KEYS: { GRAPH_NAME: the name of the graph file SUMMARY: dimension of embedding, HIDDENCHANNELS: the hidden channel of encoder LEARNING_RATE: learning rate, BATCH_SIZE: batch size of each batch, NUM_EPOCH: number of epoch to be trained, CUDA: use GPU }] feature ([np.array]): [the node features] Returns: [np.array]: [numpy array format of the DGI embedding] """ g = io.loadmat(osp.join(fp, 'interim', 'graph', PARAMS['GRAPH_NAME'])) N = g['N'] p_cate = feature post_indx = g['post_indx'] edge_idx, x = from_scipy_sparse_matrix(N) x = x.view(-1, 1).float() feature = np.zeros((x.shape[0], p_cate.shape[1])) feature[post_indx, :] = p_cate x = torch.cat([x, torch.FloatTensor(feature)], 1) data = Data(x=x, edge_index=edge_idx) if PARAMS['CUDA']: device = 'cuda' if torch.cuda.is_available() else 'cpu' else: device = 'cpu' data = data.to(device) model = DeepGraphInfomax( hidden_channels=PARAMS['HIDDEN_CHANNELS'], encoder=Encoder(data.x.shape[1], PARAMS['SUMMARY']), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=PARAMS['LEARNING_RATE']) def train(): model.train() optimizer.zero_grad() pos_z, neg_z, summary = model(data.x, data.edge_index) loss = model.loss(pos_z, neg_z, summary) loss.backward() optimizer.step() return loss.item() losses = [] for epoch in range(1, PARAMS['NUM_EPOCH'] + 1): loss = train() losses.append(loss) print('Epoch: {:03d}, Loss: {:.4f}'.format(epoch, loss)) model.eval() with torch.no_grad(): z, _, _ = model(data.x, data.edge_index) if not os.path.exists(os.path.join(fp, 'processed', 'infomax')): os.mkdir(os.path.join(fp, 'processed', 'infomax')) with open( osp.join(fp, 'processed', 'infomax', PARAMS['EMBEDDING_NAME'] + 'log.json'), 'w') as f: json.dump({'loss': losses}, f) z = z.detach().cpu().numpy()[post_indx.reshape(-1, ), :] np.save(osp.join(fp, 'processed', 'infomax', PARAMS['EMBEDDING_NAME']), z) print('embedding infomax created') return z
def __init__(self, inp_dim, out_dim, device): self.encoder = DGIEncoderNet(inp_dim, out_dim) self.dgi = DeepGraphInfomax(out_dim, encoder=self.encoder, summary=self.readout, corruption=self.corrupt) self.dgi = self.dgi.to(device) self.optimizer = torch.optim.Adam(self.dgi.parameters())
def train(): # get the parameters args = get_args() print(args.domain) # decide the device device = torch.device('cuda:2' if torch.cuda.is_available() and args.cuda else 'cpu') # load dataset if args.domain == 'Cora': dataset = Planetoid(root='/home/amax/xsx/data/gnn_datas/Cora', name='Cora', transform=T.NormalizeFeatures()) elif args.domain == 'CiteSeer': dataset = Planetoid(root='/home/amax/xsx/data/gnn_datas/CiteSeer', name='CiteSeer', transform=T.NormalizeFeatures()) elif args.domain == 'PubMed': dataset = Planetoid(root='/home/amax/xsx/data/gnn_datas/PubMed', name='PubMed', transform=T.NormalizeFeatures()) elif args.domain == 'DBLP': dataset = DBLP(root='/home/amax/xsx/data/gnn_datas/DBLP', name='DBLP') elif args.domain == 'Cora-ML': dataset = CoraML(root='/home/amax/xsx/data/gnn_datas/Cora_ML', name='Cora_ML') elif args.domain == 'CS': dataset = Coauthor(root='/home/amax/xsx/data/gnn_datas/Coauthor/CS', name='CS') elif args.domain == 'Physics': dataset = Coauthor(root='/home/amax/xsx/data/gnn_datas/Coauthor/Physics', name='Physics') elif args.domain == 'Computers': dataset = Amazon(root='/home/amax/xsx/data/gnn_datas/Amazon/Computers', name='Computers') elif args.domain == 'Photo': dataset = Amazon(root='/home/amax/xsx/data/gnn_datas/Amazon/Photo', name='Photo') else: dataset = None if dataset is None: pdb.set_trace() data = dataset[0].to(device) # create the model and optimizer model = DeepGraphInfomax(hidden_channels=args.hidden_dim, encoder=Encoder(dataset.num_features, args.hidden_dim), summary=lambda z, *args, **kwargs: z.mean(dim=0), corruption=corruption).to(device) optimizer = Adam(model.parameters(), lr=args.lr) # the information which need to be recorded start_time = time.time() bad_counter = 0 best_epoch = 0 least_loss = float("inf") best_model = None # beging training for epoch in range(args.epochs): # the steps of training model.train() optimizer.zero_grad() pos_z, neg_z, summary = model(data.x, data.edge_index) loss = model.loss(pos_z, neg_z, summary) current_loss = loss.item() loss.backward() optimizer.step() # save the model if it access the minimum loss in current epoch if current_loss < least_loss: least_loss = current_loss best_epoch = epoch + 1 best_model = copy.deepcopy(model) bad_counter = 0 else: bad_counter += 1 # early stop if bad_counter >= args.patience: break print("Optimization Finished!") used_time = time.time() - start_time print("Total epochs: {:2d}".format(best_epoch + 100)) print("Best epochs: {:2d}".format(best_epoch)) # train a classification model node_classification(best_model, data, args, device, int(dataset.num_classes)) print("Total time elapsed: {:.2f}s".format(used_time))
def run_test(self, dataset_getter, logger, other=None): """ This function returns the training and test accuracy. DO NOT USE THE TEST FOR TRAINING OR EARLY STOPPING! :return: (training accuracy, test accuracy) """ dataset_class = self.model_config.dataset # dataset_class() dataset = dataset_class() device = self.model_config.device shuffle = self.model_config[ 'shuffle'] if 'shuffle' in self.model_config else True train_loader, val_loader = dataset_getter.get_train_val( dataset, self.model_config['batch_size'], shuffle=shuffle) test_loader = dataset_getter.get_test(dataset, self.model_config['batch_size'], shuffle=shuffle) def corruption(x, edge_index, dataset): # Ignore parameters, we do not want to shuffle the graph features as in single graph tasks rand_idx = random.randint(0, len(dataset) - 1) return dataset[rand_idx].x.to( device), dataset[rand_idx].edge_index.to(device) model = DeepGraphInfomax( hidden_channels=self.model_config['dim_embedding'], encoder=DGIEncoder(dataset._dim_features, self.model_config['dim_embedding'], self.model_config), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=lambda x, edge_index: corruption( x, edge_index, train_loader.dataset)).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=self.model_config['learning_rate']) loss_fun = nn.BCELoss() # Train for epoch in range(1, self.model_config['DGI_epochs'] + 1): loss = train(model, train_loader.dataset, device, loss_fun, optimizer, self.model_config['training_strategy']) print('Epoch: {:03d}, Loss: {:.4f}'.format(epoch, loss)) # Compute embeddings train_emb, train_y = eval(model, train_loader.dataset, device) valid_emb, valid_y = eval(model, val_loader.dataset, device) test_emb, test_y = eval(model, test_loader.dataset, device) train_emb = train_emb.detach().cpu().numpy() train_y = train_y.detach().cpu().numpy() valid_emb = valid_emb.detach().cpu().numpy() valid_y = valid_y.detach().cpu().numpy() test_emb = test_emb.detach().cpu().numpy() test_y = test_y.detach().cpu().numpy() # Standardize embeddings mean = train_emb.mean(axis=0) std = train_emb.std(axis=0) train_emb = (train_emb - mean) / std valid_emb = (valid_emb - mean) / std test_emb = (test_emb - mean) / std train_score, validation_score, test_score = run_regression( train_emb, train_y, valid_emb, valid_y, test_emb, test_y) print( f'MICRO F1 TRAIN SCORE: {train_score} \t VALID SCORE: {validation_score} \t TEST SCORE: {test_score}' ) return train_score, test_score return train_acc, test_acc
def main(eps, hidden_size): class Encoder_GIN(nn.Module): def __init__(self, in_channels, hidden_channels): super(Encoder_GIN, self).__init__() nn1 = Sequential(Linear(in_channels, hidden_size)) self.conv1 = GINConv(nn1, eps=eps) self.prelu = nn.PReLU(hidden_channels) def forward(self, x, edge_index): x = self.conv1(x, edge_index) x = self.prelu(x) return x dataset = 'Citeseer' path = osp.join(osp.dirname(osp.realpath(sys.argv[0])), '..', 'data', dataset) dataset = Planetoid(path, dataset) class Encoder(nn.Module): def __init__(self, in_channels, hidden_channels): super(Encoder, self).__init__() self.conv = GCNConv(in_channels, hidden_channels, cached=True) self.prelu = nn.PReLU(hidden_channels) def forward(self, x, edge_index): x = self.conv(x, edge_index) x = self.prelu(x) return x def corruption(x, edge_index): return x[torch.randperm(x.size(0))], edge_index device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # model = DeepGraphInfomax( # hidden_channels=512, encoder=Encoder(dataset.num_features, 512), # summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), # corruption=corruption).to(device) model = DeepGraphInfomax( hidden_channels=512, encoder=Encoder_GIN(dataset.num_features, 512), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption).to(device) data = dataset[0].to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) def train(): model.train() optimizer.zero_grad() pos_z, neg_z, summary = model(data.x, data.edge_index) loss = model.loss(pos_z, neg_z, summary) loss.backward() optimizer.step() return loss.item() def test(): model.eval() z, _, _ = model(data.x, data.edge_index) acc = model.test(z[data.train_mask], data.y[data.train_mask], z[data.test_mask], data.y[data.test_mask], max_iter=150) return acc for epoch in range(1, 301): loss = train() #print('Epoch: {:03d}, Loss: {:.4f}'.format(epoch, loss)) acc = test() # print('Accuracy: {:.4f}'.format(acc)) return acc
savename_suffix = f'infomax_react_{encoder_model}' if encoder_model != 'graph_finetuned': savename_suffix += f'-gr{gr}' model_dir = f'./models_react/20k_data/{savename_suffix}' writer = SummaryWriter(f'logs/20k_data/{savename_suffix}') save_dir = os.path.join(model_dir, 'weights') os.makedirs(save_dir, exist_ok=True) model = DeepGraphInfomax( hidden_channels=512, encoder=Encoder( PARAM_FILE, LANG_FILE, WEIGHT_FILE, device=device, # TODO Is this the best spot? in_channels=256, # Latent space size hidden_channels=512, batch_mode='Packed' ), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption ).to(device) optimizer = OPTIMIZER_FACTORY['adam'](model.parameters(), lr=1e-4) scheduler = None # Loads the entire data from the processed files loader = load_data( PROCESSED_FILE, LANG_FILE, max_total_bundle=max_total_bundle )
def main_model_dgi(data, hidden, if_all=False): torch.backends.cudnn.deterministic = True device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = DeepGraphInfomax( hidden_channels=hidden, encoder=Encoder(hidden, data), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption) data.split_train_valid() model = model.to(device) data = data.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) best_acc_valid = 0 for epoch in range(10): model.train() optimizer.zero_grad() pos_z, neg_z, summary = model(data.x, data.edge_index) lr = LogisticRegression().fit(pos_z[data.mask_train].detach().cpu().numpy().reshape(-1, hidden), data.y[data.mask_train].cpu().numpy()) valid_pred = lr.predict(pos_z[data.mask_valid].detach().cpu().numpy().reshape(-1, hidden)) acc_valid = accuracy_score(data.y[data.mask_valid].cpu().numpy(), valid_pred) if acc_valid > best_acc_valid: best_acc_valid = acc_valid result = pos_z loss = model.loss(pos_z.to(device), neg_z.to(device), summary.to(device)) loss.backward() optimizer.step() lr = LogisticRegression().fit(result[data.mask_train].detach().cpu().numpy().reshape(-1, hidden), data.y[data.mask_train].cpu().numpy()) train_pred = lr.predict(result[data.mask_train].detach().cpu().numpy().reshape(-1, hidden)) all_pred = lr.predict(result.detach().cpu().numpy().reshape(-1, hidden)) if if_all: return Result( result=torch.tensor(np.eye(data.num_class)[all_pred]).float().cpu(), loss_train=-1, loss_valid=-1, acc_train=accuracy_score(data.y[data.mask_train].cpu().numpy(), train_pred), acc_valid=best_acc_valid, epoch=10, ) else: return Result( result=all_pred[data.mask_test], loss_train=-1, loss_valid=-1, acc_train=accuracy_score(data.y[data.mask_train].cpu().numpy(), train_pred), acc_valid=best_acc_valid, epoch=10, )
def test(): model.eval() z, _, _ = model(data.x, data.edge_index) acc = model.test(z[data.train_mask], data.y[data.train_mask], z[data.test_mask], data.y[data.test_mask], max_iter=150) return acc test_acc = [] best_acc = 0 for _ in range(5): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = DeepGraphInfomax( hidden_channels=512, encoder=Encoder(dataset.num_features, 512), summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)), corruption=corruption).to(device) data = dataset[0].to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) for epoch in range(1, 301): loss = train() # print('Epoch: {:03d}, Loss: {:.4f}'.format(epoch, loss)) acc = test() test_acc.append(acc) print('Accuracy: {:.4f}'.format(acc)) if best_acc < acc: best_acc = acc if not os.path.isdir('../checkpoint'): os.makedirs('../checkpoint') torch.save(model.state_dict(), os.path.join('../checkpoint', '{}.pth'.format(args.dataset))) test_acc = np.array(test_acc)