Updating to PyTorch 0.4.0

Author: dasguptar

main.py

@@ -47,6 +47,7 @@ def main():
     logger.addHandler(ch)
     # argument validation
     args.cuda = args.cuda and torch.cuda.is_available()
+    device = torch.device("cuda:0" if args.cuda else "cpu")
     if args.sparse and args.wd != 0:
         logger.error('Sparsity and weight decay are incompatible, pick one!')
         exit()
@@ -111,18 +112,6 @@ def main():
         args.sparse,
         args.freeze_embed)
     criterion = nn.KLDivLoss()
-    if args.cuda:
-        model.cuda(), criterion.cuda()
-    if args.optim == 'adam':
-        optimizer = optim.Adam(filter(lambda p: p.requires_grad,
-                                      model.parameters()), lr=args.lr, weight_decay=args.wd)
-    elif args.optim == 'adagrad':
-        optimizer = optim.Adagrad(filter(lambda p: p.requires_grad,
-                                         model.parameters()), lr=args.lr, weight_decay=args.wd)
-    elif args.optim == 'sgd':
-        optimizer = optim.SGD(filter(lambda p: p.requires_grad,
-                                     model.parameters()), lr=args.lr, weight_decay=args.wd)
-    metrics = Metrics(args.num_classes)
 
     # for words common to dataset vocab and GLOVE, use GLOVE vectors
     # for other words in dataset vocab, use random normal vectors
@@ -134,7 +123,8 @@ def main():
         glove_vocab, glove_emb = utils.load_word_vectors(
             os.path.join(args.glove, 'glove.840B.300d'))
         logger.debug('==> GLOVE vocabulary size: %d ' % glove_vocab.size())
-        emb = torch.Tensor(vocab.size(), glove_emb.size(1)).normal_(-0.05, 0.05)
+        emb = torch.zeros(vocab.size(), glove_emb.size(1), dtype=torch.float, device=device)
+        emb.normal_(0, 0.05)
         # zero out the embeddings for padding and other special words if they are absent in vocab
         for idx, item in enumerate([Constants.PAD_WORD, Constants.UNK_WORD,
                                     Constants.BOS_WORD, Constants.EOS_WORD]):
@@ -144,12 +134,22 @@ def main():
                 emb[vocab.getIndex(word)] = glove_emb[glove_vocab.getIndex(word)]
         torch.save(emb, emb_file)
     # plug these into embedding matrix inside model
-    if args.cuda:
-        emb = emb.cuda()
-    model.emb.weight.data.copy_(emb)
+    model.emb.weight.copy_(emb)
+
+    model.to(device), criterion.to(device)
+    if args.optim == 'adam':
+        optimizer = optim.Adam(filter(lambda p: p.requires_grad,
+                                      model.parameters()), lr=args.lr, weight_decay=args.wd)
+    elif args.optim == 'adagrad':
+        optimizer = optim.Adagrad(filter(lambda p: p.requires_grad,
+                                         model.parameters()), lr=args.lr, weight_decay=args.wd)
+    elif args.optim == 'sgd':
+        optimizer = optim.SGD(filter(lambda p: p.requires_grad,
+                                     model.parameters()), lr=args.lr, weight_decay=args.wd)
+    metrics = Metrics(args.num_classes)
 
     # create trainer object for training and testing
-    trainer = Trainer(args, model, criterion, optimizer)
+    trainer = Trainer(args, model, criterion, optimizer, device)
 
     best = -float('inf')
     for epoch in range(args.epochs):

requirements.txt

@@ -1,2 +1,2 @@
-http://download.pytorch.org/whl/cpu/torch-0.3.1-cp36-cp36m-linux_x86_64.whl
+http://download.pytorch.org/whl/cpu/torch-0.4.0-cp36-cp36m-linux_x86_64.whl
 tqdm

treelstm/dataset.py

@@ -44,7 +44,7 @@ def read_sentences(self, filename):
 
     def read_sentence(self, line):
         indices = self.vocab.convertToIdx(line.split(), Constants.UNK_WORD)
-        return torch.LongTensor(indices)
+        return torch.tensor(indices, dtype=torch.long, device='cpu')
 
     def read_trees(self, filename):
         with open(filename, 'r') as f:
@@ -82,5 +82,5 @@ def read_tree(self, line):
     def read_labels(self, filename):
         with open(filename, 'r') as f:
             labels = list(map(lambda x: float(x), f.readlines()))
-            labels = torch.Tensor(labels)
+            labels = torch.tensor(labels, dtype=torch.float, device='cpu')
         return labels

treelstm/model.py

@@ -1,7 +1,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.autograd import Variable as Var
 
 from . import Constants
 
@@ -39,8 +38,8 @@ def forward(self, tree, inputs):
             self.forward(tree.children[idx], inputs)
 
         if tree.num_children == 0:
-            child_c = Var(inputs[0].data.new(1, self.mem_dim).fill_(0.))
-            child_h = Var(inputs[0].data.new(1, self.mem_dim).fill_(0.))
+            child_c = inputs[0].detach().new(1, self.mem_dim).fill_(0.).requires_grad_()
+            child_h = inputs[0].detach().new(1, self.mem_dim).fill_(0.).requires_grad_()
         else:
             child_c, child_h = zip(* map(lambda x: x.state, tree.children))
             child_c, child_h = torch.cat(child_c, dim=0), torch.cat(child_h, dim=0)

treelstm/trainer.py

@@ -1,36 +1,34 @@
 from tqdm import tqdm
 
 import torch
-from torch.autograd import Variable as Var
 
 from . import utils
 
 
 class Trainer(object):
-    def __init__(self, args, model, criterion, optimizer):
+    def __init__(self, args, model, criterion, optimizer, device):
         super(Trainer, self).__init__()
         self.args = args
         self.model = model
         self.criterion = criterion
         self.optimizer = optimizer
+        self.device = device
         self.epoch = 0
 
     # helper function for training
     def train(self, dataset):
         self.model.train()
         self.optimizer.zero_grad()
         total_loss = 0.0
-        indices = torch.randperm(len(dataset))
+        indices = torch.randperm(len(dataset), dtype=torch.long, device='cpu')
         for idx in tqdm(range(len(dataset)), desc='Training epoch ' + str(self.epoch + 1) + ''):
-            ltree, lsent, rtree, rsent, label = dataset[indices[idx]]
-            linput, rinput = Var(lsent), Var(rsent)
-            target = Var(utils.map_label_to_target(label, dataset.num_classes))
-            if self.args.cuda:
-                linput, rinput = linput.cuda(), rinput.cuda()
-                target = target.cuda()
+            ltree, linput, rtree, rinput, label = dataset[indices[idx]]
+            target = utils.map_label_to_target(label, dataset.num_classes)
+            linput, rinput = linput.to(self.device), rinput.to(self.device)
+            target = target.to(self.device)
             output = self.model(ltree, linput, rtree, rinput)
             loss = self.criterion(output, target)
-            total_loss += loss.data[0]
+            total_loss += loss.item()
             loss.backward()
             if idx % self.args.batchsize == 0 and idx > 0:
                 self.optimizer.step()
@@ -41,19 +39,18 @@ def train(self, dataset):
     # helper function for testing
     def test(self, dataset):
         self.model.eval()
-        total_loss = 0
-        predictions = torch.zeros(len(dataset))
-        indices = torch.arange(1, dataset.num_classes + 1)
-        for idx in tqdm(range(len(dataset)), desc='Testing epoch  ' + str(self.epoch) + ''):
-            ltree, lsent, rtree, rsent, label = dataset[idx]
-            linput, rinput = Var(lsent, volatile=True), Var(rsent, volatile=True)
-            target = Var(utils.map_label_to_target(label, dataset.num_classes), volatile=True)
-            if self.args.cuda:
-                linput, rinput = linput.cuda(), rinput.cuda()
-                target = target.cuda()
-            output = self.model(ltree, linput, rtree, rinput)
-            loss = self.criterion(output, target)
-            total_loss += loss.data[0]
-            output = output.data.squeeze().cpu()
-            predictions[idx] = torch.dot(indices, torch.exp(output))
+        with torch.no_grad():
+            total_loss = 0.0
+            predictions = torch.zeros(len(dataset), dtype=torch.float, device='cpu')
+            indices = torch.arange(1, dataset.num_classes + 1, dtype=torch.float, device='cpu')
+            for idx in tqdm(range(len(dataset)), desc='Testing epoch  ' + str(self.epoch) + ''):
+                ltree, linput, rtree, rinput, label = dataset[idx]
+                target = utils.map_label_to_target(label, dataset.num_classes)
+                linput, rinput = linput.to(self.device), rinput.to(self.device)
+                target = target.to(self.device)
+                output = self.model(ltree, linput, rtree, rinput)
+                loss = self.criterion(output, target)
+                total_loss += loss.item()
+                output = output.squeeze().to('cpu')
+                predictions[idx] = torch.dot(indices, torch.exp(output))
         return total_loss / len(dataset), predictions

treelstm/utils.py

@@ -26,13 +26,14 @@ def load_word_vectors(path):
         contents = f.readline().rstrip('\n').split(' ')
         dim = len(contents[1:])
     words = [None] * (count)
-    vectors = torch.zeros(count, dim)
+    vectors = torch.zeros(count, dim, dtype=torch.float, device='cpu')
     with open(path + '.txt', 'r', encoding='utf8', errors='ignore') as f:
         idx = 0
         for line in f:
             contents = line.rstrip('\n').split(' ')
             words[idx] = contents[0]
-            vectors[idx] = torch.Tensor(list(map(float, contents[1:])))
+            values = list(map(float, contents[1:]))
+            vectors[idx] = torch.tensor(values, dtype=torch.float, device='cpu')
             idx += 1
     with open(path + '.vocab', 'w', encoding='utf8', errors='ignore') as f:
         for word in words:
@@ -57,12 +58,12 @@ def build_vocab(filenames, vocabfile):
 
 # mapping from scalar to vector
 def map_label_to_target(label, num_classes):
-    target = torch.zeros(1, num_classes)
+    target = torch.zeros(1, num_classes, dtype=torch.float, device='cpu')
     ceil = int(math.ceil(label))
     floor = int(math.floor(label))
     if ceil == floor:
-        target[0][floor-1] = 1
+        target[0, floor-1] = 1
     else:
-        target[0][floor-1] = ceil - label
-        target[0][ceil-1] = label - floor
+        target[0, floor-1] = ceil - label
+        target[0, ceil-1] = label - floor
     return target