dmlc--dgl
S22
b04cf65b2b
It seems that there is a bug in Sparseadam in PyTorch, and it can be temporarily fixed by: optimizer = optim.SparseAdam(list(self.skip_gram_model.parameters()), lr=self.initial_lr) Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>
82 行
3.6 KiB
Python
82 行
3.6 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
|
|
from download import AminerDataset, CustomDataset
|
|
|
|
|
|
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()
|