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()
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.
adjs = [adj.to(device) for adj in adjs]
optimizer.zero_grad()
pos_z, neg_z, summary = model(x[n_id], adjs)
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 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
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
)
# Should load intermediate too, but for now let's just grab those
# where we actually finished the training.
last_model = os.path.join(save_dir, f'saved_model_epoch_{epochs}.pt')
if os.path.exists(last_model):
model.load_state_dict(torch.load(last_model, map_location=device))
else:
for epoch in range(epochs):
train(
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,
)
