wehub-resource-sync
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import argparse
|
|
|
|
import torch
|
|
import torch.optim as optim
|
|
from download import AminerDataset, CustomDataset
|
|
from model import SkipGramModel
|
|
|
|
from reading_data import DataReader, Metapath2vecDataset
|
|
from torch.utils.data import DataLoader
|
|
|
|
from tqdm import tqdm
|
|
|
|
|
|
class Metapath2VecTrainer:
|
|
def __init__(self, args):
|
|
if args.aminer:
|
|
dataset = AminerDataset(args.path)
|
|
else:
|
|
dataset = CustomDataset(args.path)
|
|
self.data = DataReader(dataset, args.min_count, args.care_type)
|
|
dataset = Metapath2vecDataset(self.data, args.window_size)
|
|
self.dataloader = DataLoader(
|
|
dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=True,
|
|
num_workers=args.num_workers,
|
|
collate_fn=dataset.collate,
|
|
)
|
|
|
|
self.output_file_name = args.output_file
|
|
self.emb_size = len(self.data.word2id)
|
|
self.emb_dimension = args.dim
|
|
self.batch_size = args.batch_size
|
|
self.iterations = args.iterations
|
|
self.initial_lr = args.initial_lr
|
|
self.skip_gram_model = SkipGramModel(self.emb_size, self.emb_dimension)
|
|
|
|
self.use_cuda = torch.cuda.is_available()
|
|
self.device = torch.device("cuda" if self.use_cuda else "cpu")
|
|
if self.use_cuda:
|
|
self.skip_gram_model.cuda()
|
|
|
|
def train(self):
|
|
optimizer = optim.SparseAdam(
|
|
list(self.skip_gram_model.parameters()), lr=self.initial_lr
|
|
)
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
|
optimizer, len(self.dataloader)
|
|
)
|
|
|
|
for iteration in range(self.iterations):
|
|
print("\n\n\nIteration: " + str(iteration + 1))
|
|
running_loss = 0.0
|
|
for i, sample_batched in enumerate(tqdm(self.dataloader)):
|
|
if len(sample_batched[0]) > 1:
|
|
pos_u = sample_batched[0].to(self.device)
|
|
pos_v = sample_batched[1].to(self.device)
|
|
neg_v = sample_batched[2].to(self.device)
|
|
|
|
scheduler.step()
|
|
optimizer.zero_grad()
|
|
loss = self.skip_gram_model.forward(pos_u, pos_v, neg_v)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
running_loss = running_loss * 0.9 + loss.item() * 0.1
|
|
if i > 0 and i % 500 == 0:
|
|
print(" Loss: " + str(running_loss))
|
|
|
|
self.skip_gram_model.save_embedding(
|
|
self.data.id2word, self.output_file_name
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Metapath2vec")
|
|
# parser.add_argument('--input_file', type=str, help="input_file")
|
|
parser.add_argument(
|
|
"--aminer", action="store_true", help="Use AMiner dataset"
|
|
)
|
|
parser.add_argument("--path", type=str, help="input_path")
|
|
parser.add_argument("--output_file", type=str, help="output_file")
|
|
parser.add_argument(
|
|
"--dim", default=128, type=int, help="embedding dimensions"
|
|
)
|
|
parser.add_argument(
|
|
"--window_size", default=7, type=int, help="context window size"
|
|
)
|
|
parser.add_argument("--iterations", default=5, type=int, help="iterations")
|
|
parser.add_argument("--batch_size", default=50, type=int, help="batch size")
|
|
parser.add_argument(
|
|
"--care_type",
|
|
default=0,
|
|
type=int,
|
|
help="if 1, heterogeneous negative sampling, else normal negative sampling",
|
|
)
|
|
parser.add_argument(
|
|
"--initial_lr", default=0.025, type=float, help="learning rate"
|
|
)
|
|
parser.add_argument("--min_count", default=5, type=int, help="min count")
|
|
parser.add_argument(
|
|
"--num_workers", default=16, type=int, help="number of workers"
|
|
)
|
|
args = parser.parse_args()
|
|
m2v = Metapath2VecTrainer(args)
|
|
m2v.train()
|