diff --git a/embed.py b/embed.py index 97f07d0..ae5a283 100644 --- a/embed.py +++ b/embed.py @@ -113,7 +113,7 @@ def control(queue, log, types, data, fout, distfn, nepochs, processes): parser.add_argument('-debug', help='Print debug output', action='store_true', default=False) opt = parser.parse_args() - th.set_default_tensor_type('torch.DoubleTensor') + th.set_default_tensor_type('torch.FloatTensor') if opt.debug: log_level = logging.DEBUG else: diff --git a/example.sh b/example.sh index 8d24fad..c3100d8 100755 --- a/example.sh +++ b/example.sh @@ -13,9 +13,9 @@ export OMP_NUM_THREADS=1 python3 embed.py \ -dim 5 \ -lr 0.3 \ - -epochs 200 \ + -epochs 300 \ -negs 50 \ - -burnin 10 \ + -burnin 20 \ -nproc "${NTHREADS}" \ -distfn poincare \ -dset wordnet/mammal_closure.tsv \ diff --git a/model.py b/model.py index 04ff39a..708af22 100644 --- a/model.py +++ b/model.py @@ -23,7 +23,6 @@ def __init__(self, eps=eps): self.eps = eps def forward(self, x): - self.save_for_backward(x) self.z = th.sqrt(x * x - 1) return th.log(x + self.z) @@ -130,7 +129,7 @@ def loss(self, preds, targets, weight=None, size_average=True): class GraphDataset(Dataset): _ntries = 10 - _dampening = 0.75 + _dampening = 1 def __init__(self, idx, objects, nnegs, unigram_size=1e8): print('Indexing data') @@ -174,7 +173,10 @@ def __getitem__(self, i): t, h, _ = self.idx[i] negs = set() ntries = 0 - while ntries < self.max_tries and len(negs) < self.nnegs: + nnegs = self.nnegs + if self.burnin: + nnegs *= 0.1 + while ntries < self.max_tries and len(negs) < nnegs: if self.burnin: n = randint(0, len(self.unigram_table)) n = int(self.unigram_table[n]) @@ -186,7 +188,7 @@ def __getitem__(self, i): if len(negs) == 0: negs.add(t) ix = [t, h] + list(negs) - while len(ix) < self.nnegs + 2: + while len(ix) < nnegs + 2: ix.append(ix[randint(2, len(ix))]) return th.LongTensor(ix).view(1, len(ix)), th.zeros(1).long() diff --git a/rsgd.py b/rsgd.py index 57ef40b..7226047 100644 --- a/rsgd.py +++ b/rsgd.py @@ -9,7 +9,7 @@ import torch as th from torch.optim.optimizer import Optimizer, required -spten_t = th.sparse.DoubleTensor +spten_t = th.sparse.FloatTensor def poincare_grad(p, d_p): diff --git a/train.py b/train.py index 82e7c57..2f7992f 100644 --- a/train.py +++ b/train.py @@ -11,7 +11,7 @@ from torch.utils.data import DataLoader import gc -_lr_multiplier = 0.1 +_lr_multiplier = 0.01 def train_mp(model, data, optimizer, opt, log, rank, queue):