dmlc--dgl
ziqiaomeng
ddeb86f91c
* Prepare for metapath sampler This file is just for reviewing metapath sampling algorithm (Python version). * Delete metapath_sampler * Prepare for metapath sampler This file is just for reviewing metapath sampling algorithm (Python code). * Add files via upload * Create metapath2vec.md * Add files via upload * Delete data_handler.py * Delete word2vec.py * Delete word_train.py * Add files via upload Metapath2vec implementations. Metapath2vec++ needs negative sampler optimization. * Delete shuffle_training.py * Delete test.py * Add files via upload * Delete sampler.py * Delete metapath_sampler.md * Add files via upload * Update and rename shuffle_training.py to metapath2vec.py * Update reading_data.py * Update metapath2vec.md * Update metapath2vec.md * Update metapath2vec.md * Update metapath2vec.md * Update metapath2vec.md * Create label 2 * Delete label 2 * Create testing.md * Add files via upload * Create sample.md * Add files via upload * Delete sampler.py * Add files via upload * Delete googlescholar.8area.author.label.txt * Delete googlescholar.8area.venue.label.txt * Delete testing.md * Delete id_author.txt * Delete id_conf.txt * Delete paper.txt * Delete paper_author.txt * Delete paper_conf.txt * Delete sample.md * Delete sampler.py * Add files via upload * Add files via upload * Add files via upload * Delete reading_data.py * Add files via upload * Add files via upload * Delete metapath2vec.py * Add files via upload * Rename shuffle_training.py to metapath2vec.py * Update metapath2vec.md * Delete reading_data.py * add comments and remov e commented codes
76 行
3.3 KiB
Python
76 行
3.3 KiB
Python
import torch
|
|
import argparse
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader
|
|
|
|
from tqdm import tqdm
|
|
|
|
from reading_data import DataReader, Metapath2vecDataset
|
|
from model import SkipGramModel
|
|
|
|
|
|
class Metapath2VecTrainer:
|
|
def __init__(self, args):
|
|
self.data = DataReader(args.download, 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):
|
|
|
|
for iteration in range(self.iterations):
|
|
print("\n\n\nIteration: " + str(iteration + 1))
|
|
optimizer = optim.SparseAdam(self.skip_gram_model.parameters(), lr=self.initial_lr)
|
|
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, len(self.dataloader))
|
|
|
|
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('--download', type=str, help="download_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()
|