wehub-resource-sync
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
from statistics import mean
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
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)
|