dmlc--dgl
Tingzhang Zhao
e58eeebfb6
* Update README.md Add description of RECT * [Example] Add implementation of RECT [Example] Add implementation of RECT * Update classify.py Modify the class names and the function names mentioned above * Update main.py Modify the function names mentioned above * Update label_utils.py Adjust the comments * Update README.md Add the github information Co-authored-by: Tianjun Xiao <xiaotj1990327@gmail.com>
49 行
2.0 KiB
Python
49 行
2.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from statistics import mean
|
|
|
|
class LogisticRegressionClassifier(nn.Module):
|
|
''' Define a logistic regression classifier to evaluate the quality of embedding results
|
|
'''
|
|
def __init__(self, nfeat, nclass):
|
|
super(LogisticRegressionClassifier, self).__init__()
|
|
self.lrc = nn.Linear(nfeat, nclass)
|
|
|
|
def forward(self, x):
|
|
preds = self.lrc(x)
|
|
return preds
|
|
|
|
def _evaluate(model, features, labels, test_mask):
|
|
model.eval()
|
|
with torch.no_grad():
|
|
logits = model(features)
|
|
logits = logits[test_mask]
|
|
labels = labels[test_mask]
|
|
_, indices = torch.max(logits, dim=1)
|
|
correct = torch.sum(indices == labels)
|
|
return correct.item() * 1.0 / len(labels)
|
|
|
|
def _train_test_with_lrc(model, features, labels, train_mask, test_mask):
|
|
''' Under the pre-defined balanced train/test label setting, train a lrc to evaluate the embedding results. '''
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=0.2, weight_decay=5e-06)
|
|
for _ in range(100):
|
|
model.train()
|
|
optimizer.zero_grad()
|
|
output = model(features)
|
|
loss_train = F.cross_entropy(output[train_mask], labels[train_mask])
|
|
loss_train.backward()
|
|
optimizer.step()
|
|
return _evaluate(model=model, features=features, labels=labels, test_mask=test_mask)
|
|
|
|
def evaluate_embeds(features, labels, train_mask, test_mask, n_classes, cuda, test_times=10):
|
|
print("Training a logistic regression classifier with the pre-defined train/test split setting ...")
|
|
res_list = []
|
|
for _ in range(test_times):
|
|
model = LogisticRegressionClassifier(nfeat=features.shape[1], nclass=n_classes)
|
|
if cuda:
|
|
model.cuda()
|
|
res = _train_test_with_lrc(model=model, features=features, labels=labels, train_mask=train_mask, test_mask=test_mask)
|
|
res_list.append(res)
|
|
return mean(res_list)
|