Use torchtext for MP-CNN (#76)

* Add SICK torchtext Dataset

* SICK dataset - torchtext postprocess into class probs

* Update model, driver, trainer, evaluator for SICK for torchtext

* MP-CNN: Fix bugs that prevent SICK from running on gpu 0

* MP-CNN: make SICK dataset w/ torchtext GPU-agnostic

* MP-CNN: support sparse features / idf overlap with torchtext

* Add MSRVID dataset with torchtext and update MP-CNN code to use it

* MP-CNN: Make torchtext deterministic by setting python random seed

* SICK and MSRVID datasets - add pair id for debug and build test vocab

* MP-CNN: Update readme to address potential module not found error

* MP-CNN: address review comments, can run on cpu

datasets/idf_utils.py

@@ -0,0 +1,52 @@
+"""
+Utilities to compute IDF scores.
+"""
+from collections import defaultdict
+
+import nltk
+nltk.download('stopwords', quiet=True)
+from nltk.corpus import stopwords
+import numpy as np
+
+
+def get_pairwise_word_to_doc_freq(sent_list_1, sent_list_2):
+    """
+    Get pairwise word to document frequency.
+    For index i, if sentence i in sent_list_1 and sentence i in sent_list_2 both
+    container word w, then w is counted only once.
+    Returns a dictionary mapping words to number of sentence pairs the word appears in.
+    """
+    word_to_doc_cnt = defaultdict(int)
+
+    for s1, s2 in zip(sent_list_1, sent_list_2):
+        unique_tokens = set(s1) | set(s2)
+        for t in unique_tokens:
+            word_to_doc_cnt[t] += 1
+
+    return word_to_doc_cnt
+
+
+def get_pairwise_overlap_features(sent_list_1, sent_list_2, word_to_doc_cnt):
+    """
+    Get overlap, idf weighted overlap, overlap excluding stopwords, and idf weighted overlap excluding stopwords.
+    """
+    stoplist = set(stopwords.words('english'))
+    num_docs = len(sent_list_1)
+    overlap_feats = []
+
+    for s1, s2 in zip(sent_list_1, sent_list_2):
+        tokens_a_set, tokens_b_set = set(s1), set(s2)
+        intersect = tokens_a_set & tokens_b_set
+        overlap = len(intersect) / (len(tokens_a_set) + len(tokens_b_set))
+        idf_intersect = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect)
+        idf_weighted_overlap = idf_intersect / (len(tokens_a_set) + len(tokens_b_set))
+
+        tokens_a_set_no_stop = set(w for w in s1 if w not in stoplist)
+        tokens_b_set_no_stop = set(w for w in s2 if w not in stoplist)
+        intersect_no_stop = tokens_a_set_no_stop & tokens_b_set_no_stop
+        overlap_no_stop = len(intersect_no_stop) / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
+        idf_intersect_no_stop = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect_no_stop)
+        idf_weighted_overlap_no_stop = idf_intersect_no_stop / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
+        overlap_feats.append([overlap, idf_weighted_overlap, overlap_no_stop, idf_weighted_overlap_no_stop])
+
+    return overlap_feats

datasets/msrvid.py

@@ -0,0 +1,94 @@
+import math
+import os
+
+import numpy as np
+import torch
+from torchtext.data.dataset import Dataset
+from torchtext.data.example import Example
+from torchtext.data.field import Field
+from torchtext.data.iterator import BucketIterator
+from torchtext.data.pipeline import Pipeline
+from torchtext.vocab import Vectors
+
+from datasets.idf_utils import get_pairwise_word_to_doc_freq, get_pairwise_overlap_features
+
+
+def get_class_probs(sim, *args):
+    """
+    Convert a single label into class probabilities.
+    """
+    class_probs = np.zeros(MSRVID.NUM_CLASSES)
+    ceil, floor = math.ceil(sim), math.floor(sim)
+    if ceil == floor:
+        class_probs[floor] = 1
+    else:
+        class_probs[floor] = ceil - sim
+        class_probs[ceil] = sim - floor
+
+    return class_probs
+
+
+class MSRVID(Dataset):
+    NAME = 'msrvid'
+    NUM_CLASSES = 6
+    ID_FIELD = Field(sequential=False, use_vocab=False, batch_first=True)
+    TEXT_FIELD = Field(batch_first=True, tokenize=lambda x: x)  # tokenizer is identity since we already tokenized it to compute external features
+    EXT_FEATS_FIELD = Field(tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, tokenize=lambda x: x)
+    LABEL_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, postprocessing=Pipeline(get_class_probs))
+
+    @staticmethod
+    def sort_key(ex):
+        return len(ex.a)
+
+    def __init__(self, path):
+        """
+        Create a MSRVID dataset instance
+        """
+        fields = [('id', self.ID_FIELD), ('a', self.TEXT_FIELD), ('b', self.TEXT_FIELD), ('ext_feats', self.EXT_FEATS_FIELD), ('label', self.LABEL_FIELD)]
+
+        examples = []
+        f1 = open(os.path.join(path, 'a.txt'), 'r')
+        f2 = open(os.path.join(path, 'b.txt'), 'r')
+        id_file = open(os.path.join(path, 'id.txt'), 'r')
+        label_file = open(os.path.join(path, 'sim.txt'), 'r')
+
+        sent_list_1 = [l.rstrip('.\n').split(' ') for l in f1]
+        sent_list_2 = [l.rstrip('.\n').split(' ') for l in f2]
+
+        word_to_doc_cnt = get_pairwise_word_to_doc_freq(sent_list_1, sent_list_2)
+        overlap_feats = get_pairwise_overlap_features(sent_list_1, sent_list_2, word_to_doc_cnt)
+
+        for pair_id, l1, l2, ext_feats, label in zip(id_file, sent_list_1, sent_list_2, overlap_feats, label_file):
+            pair_id = pair_id.rstrip('.\n')
+            label = label.rstrip('.\n')
+            example = Example.fromlist([pair_id, l1, l2, ext_feats, label], fields)
+            examples.append(example)
+
+        map(lambda f: f.close(), [f1, f2, label_file])
+
+        super(MSRVID, self).__init__(examples, fields)
+
+    @classmethod
+    def splits(cls, path, train='train', test='test', **kwargs):
+        return super(MSRVID, cls).splits(path, train=train, test=test, **kwargs)
+
+    @classmethod
+    def iters(cls, path, vectors_name, vectors_cache, batch_size=64, shuffle=True, device=0, vectors=None, unk_init=torch.Tensor.zero_):
+        """
+        :param path: directory containing train, test, dev files
+        :param vectors_name: name of word vectors file
+        :param vectors_cache: path to word vectors file
+        :param batch_size: batch size
+        :param device: GPU device
+        :param vectors: custom vectors - either predefined torchtext vectors or your own custom Vector classes
+        :param unk_init: function used to generate vector for OOV words
+        :return:
+        """
+        if vectors is None:
+            vectors = Vectors(name=vectors_name, cache=vectors_cache, unk_init=unk_init)
+
+        train, test = cls.splits(path)
+
+        cls.TEXT_FIELD.build_vocab(train, test, vectors=vectors)
+
+        return BucketIterator.splits((train, test), batch_size=batch_size, repeat=False, shuffle=shuffle, device=device)

datasets/sick.py

@@ -0,0 +1,94 @@
+import math
+import os
+
+import numpy as np
+import torch
+from torchtext.data.dataset import Dataset
+from torchtext.data.example import Example
+from torchtext.data.field import Field
+from torchtext.data.iterator import BucketIterator
+from torchtext.data.pipeline import Pipeline
+from torchtext.vocab import Vectors
+
+from datasets.idf_utils import get_pairwise_word_to_doc_freq, get_pairwise_overlap_features
+
+
+def get_class_probs(sim, *args):
+    """
+    Convert a single label into class probabilities.
+    """
+    class_probs = np.zeros(SICK.NUM_CLASSES)
+    ceil, floor = math.ceil(sim), math.floor(sim)
+    if ceil == floor:
+        class_probs[floor - 1] = 1
+    else:
+        class_probs[floor - 1] = ceil - sim
+        class_probs[ceil - 1] = sim - floor
+
+    return class_probs
+
+
+class SICK(Dataset):
+    NAME = 'sick'
+    NUM_CLASSES = 5
+    ID_FIELD = Field(sequential=False, use_vocab=False, batch_first=True)
+    TEXT_FIELD = Field(batch_first=True, tokenize=lambda x: x)  # tokenizer is identity since we already tokenized it to compute external features
+    EXT_FEATS_FIELD = Field(tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, tokenize=lambda x: x)
+    LABEL_FIELD = Field(sequential=False, tensor_type=torch.FloatTensor, use_vocab=False, batch_first=True, postprocessing=Pipeline(get_class_probs))
+
+    @staticmethod
+    def sort_key(ex):
+        return len(ex.a)
+
+    def __init__(self, path):
+        """
+        Create a SICK dataset instance
+        """
+        fields = [('id', self.ID_FIELD), ('a', self.TEXT_FIELD), ('b', self.TEXT_FIELD), ('ext_feats', self.EXT_FEATS_FIELD), ('label', self.LABEL_FIELD)]
+
+        examples = []
+        f1 = open(os.path.join(path, 'a.txt'), 'r')
+        f2 = open(os.path.join(path, 'b.txt'), 'r')
+        id_file = open(os.path.join(path, 'id.txt'), 'r')
+        label_file = open(os.path.join(path, 'sim.txt'), 'r')
+
+        sent_list_1 = [l.rstrip('.\n').split(' ') for l in f1]
+        sent_list_2 = [l.rstrip('.\n').split(' ') for l in f2]
+
+        word_to_doc_cnt = get_pairwise_word_to_doc_freq(sent_list_1, sent_list_2)
+        overlap_feats = get_pairwise_overlap_features(sent_list_1, sent_list_2, word_to_doc_cnt)
+
+        for pair_id, l1, l2, ext_feats, label in zip(id_file, sent_list_1, sent_list_2, overlap_feats, label_file):
+            pair_id = pair_id.rstrip('.\n')
+            label = label.rstrip('.\n')
+            example = Example.fromlist([pair_id, l1, l2, ext_feats, label], fields)
+            examples.append(example)
+
+        map(lambda f: f.close(), [f1, f2, label_file])
+
+        super(SICK, self).__init__(examples, fields)
+
+    @classmethod
+    def splits(cls, path, train='train', validation='dev', test='test', **kwargs):
+        return super(SICK, cls).splits(path, train=train, validation=validation, test=test, **kwargs)
+
+    @classmethod
+    def iters(cls, path, vectors_name, vectors_cache, batch_size=64, shuffle=True, device=0, vectors=None, unk_init=torch.Tensor.zero_):
+        """
+        :param path: directory containing train, test, dev files
+        :param vectors_name: name of word vectors file
+        :param vectors_cache: path to word vectors file
+        :param batch_size: batch size
+        :param device: GPU device
+        :param vectors: custom vectors - either predefined torchtext vectors or your own custom Vector classes
+        :param unk_init: function used to generate vector for OOV words
+        :return:
+        """
+        if vectors is None:
+            vectors = Vectors(name=vectors_name, cache=vectors_cache, unk_init=unk_init)
+
+        train, val, test = cls.splits(path)
+
+        cls.TEXT_FIELD.build_vocab(train, val, test, vectors=vectors)
+
+        return BucketIterator.splits((train, val, test), batch_size=batch_size, repeat=False, shuffle=shuffle, device=device)

mp_cnn/README.md

@@ -20,7 +20,7 @@ Directory layout should be like this:
 │   └── GloVe/
 ```
 
-To run MP-CNN on the SICK dataset, use the following command. `--dropout 0` is for mimicking the original paper, although adding dropout can improve performance.
+To run MP-CNN on the SICK dataset, use the following command. `--dropout 0` is for mimicking the original paper, although adding dropout can improve performance. If you have any problems running it check the Troubleshooting section below.
 
 ```
 python main.py mpcnn.sick.model.castor --dataset sick --epochs 19 --epsilon 1e-7 --dropout 0
@@ -48,6 +48,20 @@ To see all options available, use
 python main.py --help
 ```
 
+## Troubleshooting
+
+### ModuleNotFoundError: datasets
+```
+Traceback (most recent call last):
+  File "main.py", line 9, in <module>
+    from dataset import MPCNNDatasetFactory
+  File "/u/z3tu/castorini/Castor/mp_cnn/dataset.py", line 12, in <module>
+    from datasets.sick import SICK
+ModuleNotFoundError: No module named 'datasets'
+```
+
+You need to make sure the repository root is in your `PYTHONPATH` environment variable. One way to do this is while you are in the repo root (Castor) as your current working directory, run `export PYTHONPATH=$(pwd)`.
+
 ## Optional Dependencies
 
 To optionally visualize the learning curve during training, we make use of https://github.com/lanpa/tensorboard-pytorch to connect to [TensorBoard](https://github.com/tensorflow/tensorboard). These projects require TensorFlow as a dependency, so you need to install TensorFlow before running the commands below. After these are installed, just add `--tensorboard` when running `main.py` and open TensorBoard in the browser.

mp_cnn/dataset.py

@@ -3,16 +3,14 @@ from enum import Enum
 import math
 import os
 
-import nltk
 import numpy as np
 import torch
 from torch.autograd import Variable
+import torch.nn as nn
 import torch.utils.data as data
 
-import preprocessing
-
-nltk.download('stopwords', quiet=True)
-from nltk.corpus import stopwords
+from datasets.sick import SICK
+from datasets.msrvid import MSRVID
 
 # logging setup
 import logging
@@ -26,10 +24,19 @@ ch.setFormatter(formatter)
 logger.addHandler(ch)
 
 
-class DatasetType(Enum):
-    TRAIN = 1
-    TEST = 2
-    DEV = 3
+class UnknownWorcVecCache(object):
+    """
+    Caches the first randomly generated word vector for a certain size to make it is reused.
+    """
+    cache = {}
+
+    @classmethod
+    def unk(cls, tensor):
+        size_tup = tuple(tensor.size())
+        if size_tup not in cls.cache:
+            cls.cache[size_tup] = torch.Tensor(tensor.size())
+            cls.cache[size_tup].normal_(0, 0.01)
+        return cls.cache[size_tup]
 
 
 class MPCNNDatasetFactory(object):
@@ -37,182 +44,22 @@ class MPCNNDatasetFactory(object):
     Get the corresponding Dataset class for a particular dataset.
     """
     @staticmethod
-    def get_dataset(dataset_name, word_vectors_file, batch_size, cuda, sample):
-        extra_args = {'shuffle': True}
-        dev_loader = None
-        if sample:
-            sample_indices = list(range(sample))
-            subset_random_sampler = data.sampler.SubsetRandomSampler(sample_indices)
-            extra_args['sampler'] = subset_random_sampler
-            extra_args['shuffle'] = False
+    def get_dataset(dataset_name, word_vectors_dir, word_vectors_file, batch_size, device):
         if dataset_name == 'sick':
-            train_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.TRAIN, cuda), batch_size=batch_size, **extra_args)
-            test_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.TEST, cuda), batch_size=batch_size, **extra_args)
-            dev_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.DEV, cuda), batch_size=batch_size, **extra_args)
+            dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'sick/')
+            train_loader, dev_loader, test_loader = SICK.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWorcVecCache.unk)
+            embedding_dim = SICK.TEXT_FIELD.vocab.vectors.size()
+            embedding = nn.Embedding(embedding_dim[0], embedding_dim[1])
+            embedding.weight = nn.Parameter(SICK.TEXT_FIELD.vocab.vectors)
+            return SICK, embedding, train_loader, test_loader, dev_loader
         elif dataset_name == 'msrvid':
-            train_loader = torch.utils.data.DataLoader(MSRVIDDataset(DatasetType.TRAIN, cuda), batch_size=batch_size, **extra_args)
-            test_loader = torch.utils.data.DataLoader(MSRVIDDataset(DatasetType.TEST, cuda), batch_size=batch_size, **extra_args)
+            dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'msrvid/')
+            dev_loader = None
+            train_loader, test_loader = MSRVID.iters(dataset_root, word_vectors_file, word_vectors_dir, batch_size, device=device, unk_init=UnknownWorcVecCache.unk)
+            embedding_dim = MSRVID.TEXT_FIELD.vocab.vectors.size()
+            embedding = nn.Embedding(embedding_dim[0], embedding_dim[1])
+            embedding.weight = nn.Parameter(MSRVID.TEXT_FIELD.vocab.vectors)
+            return MSRVID, embedding, train_loader, test_loader, dev_loader
         else:
             raise ValueError('{} is not a valid dataset.'.format(dataset_name))
 
-        word_index, embedding = preprocessing.get_glove_embedding(word_vectors_file, train_loader.dataset.dataset_root)
-        logger.info('Finished loading GloVe embedding for vocab in data...')
-
-        train_loader.dataset.initialize(word_index, embedding)
-        test_loader.dataset.initialize(word_index, embedding)
-        if dev_loader is not None:
-            dev_loader.dataset.initialize(word_index, embedding)
-        return train_loader, test_loader, dev_loader
-
-
-class MPCNNDataset(data.Dataset):
-    train_folder = 'train'
-    test_folder = 'test'
-    dev_folder = 'dev'
-    # subclass will override fields below
-    dataset_root = None
-    num_classes = None
-
-    def __init__(self, dataset_type, cuda):
-        if not isinstance(dataset_type, DatasetType):
-            raise ValueError('dataset_type ({}) must be of type DatasetType enum'.format(dataset_type))
-
-        if dataset_type == DatasetType.TRAIN:
-            subfolder = MPCNNDataset.train_folder
-        elif dataset_type == DatasetType.TEST:
-            subfolder = MPCNNDataset.test_folder
-        else:
-            subfolder = MPCNNDataset.dev_folder
-
-        self.dataset_dir = os.path.join(self.dataset_root, subfolder)
-        if not os.path.exists(self.dataset_dir):
-            raise RuntimeError('{} does not exist'.format(self.dataset_dir))
-
-        self.cuda = cuda
-        self.max_length = -10000
-        self.unk = torch.Tensor(300)
-        self.unk.normal_(0, 0.01)
-
-    def initialize(self, word_index, embedding):
-        """
-        Convert sentences into sentence embeddings.
-        """
-        sent_a = self._load(self.dataset_dir, 'a.txt')
-        sent_b = self._load(self.dataset_dir, 'b.txt')
-        word_to_doc_cnt = defaultdict(int)
-
-        # obtain max sentence length to use as dimension for padding to support batching
-        sent_a_tokens, sent_b_tokens = [], []
-        for i in range(len(sent_a)):
-            sa_tokens = sent_a[i].split(' ')
-            sb_tokens = sent_b[i].split(' ')
-            self.max_length = max(self.max_length, len(sa_tokens), len(sb_tokens))
-            sent_a_tokens.append(sa_tokens)
-            sent_b_tokens.append(sb_tokens)
-
-            unique_tokens = set(sa_tokens) | set(sb_tokens)
-            for t in unique_tokens:
-                word_to_doc_cnt[t] += 1
-
-        self.sentences = []
-        stoplist = set(stopwords.words('english'))
-        num_docs = len(word_to_doc_cnt)
-        for i in range(len(sent_a)):
-            sent_pair = {}
-            sent_pair['a'] = self._get_sentence_embeddings(sent_a_tokens[i], word_index, embedding)
-            sent_pair['b'] = self._get_sentence_embeddings(sent_b_tokens[i], word_index, embedding)
-
-            tokens_a_set, tokens_b_set = set(sent_a_tokens[i]), set(sent_b_tokens[i])
-            intersect = tokens_a_set & tokens_b_set
-            overlap = len(intersect) / (len(tokens_a_set) + len(tokens_b_set))
-            idf_intersect = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect)
-            idf_weighted_overlap = idf_intersect / (len(tokens_a_set) + len(tokens_b_set))
-
-            tokens_a_set_no_stop = set(w for w in sent_a_tokens[i] if w not in stoplist)
-            tokens_b_set_no_stop = set(w for w in sent_b_tokens[i] if w not in stoplist)
-            intersect_no_stop = tokens_a_set_no_stop & tokens_b_set_no_stop
-            overlap_no_stop = len(intersect_no_stop) / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
-            idf_intersect_no_stop = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect_no_stop)
-            idf_weighted_overlap_no_stop = idf_intersect_no_stop / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
-            ext_feats = torch.Tensor([overlap, idf_weighted_overlap, overlap_no_stop, idf_weighted_overlap_no_stop])
-            ext_feats = ext_feats.cuda() if self.cuda else ext_feats
-            sent_pair['ext_feats'] = ext_feats
-
-            self.sentences.append(sent_pair)
-
-        self.labels = self._load(self.dataset_dir, 'sim.txt', float)
-
-    def _load(self, dataset_dir, fname, type_converter=str):
-        data = []
-        with open(os.path.join(dataset_dir, fname), 'r') as f:
-            for line in f:
-                stripped_line = line.rstrip('.\n')
-                item = type_converter(stripped_line)
-                data.append(item)
-        return data
-
-    def _get_sentence_embeddings(self, tokens, word_index, embedding):
-        sentence_embedding = torch.zeros(300, self.max_length)
-        sentence_embedding[:, :len(tokens)].normal_(0, 1)
-        found_pos, found_emb_idx = [], []
-        for i, token in enumerate(tokens):
-            if token in word_index:
-                found_pos.append(i)
-                found_emb_idx.append(word_index[token])
-            else:
-                sentence_embedding[:, i] = self.unk
-
-        found_word_vecs = embedding(Variable(torch.LongTensor(found_emb_idx)))
-        for i, v in enumerate(found_pos):
-            sentence_embedding[:, v] = found_word_vecs[i].data
-        return sentence_embedding.cuda() if self.cuda else sentence_embedding
-
-    def __getitem__(self, idx):
-        return self.sentences[idx], self.labels[idx]
-
-    def __len__(self):
-        return len(self.labels)
-
-
-class SICKDataset(MPCNNDataset):
-
-    dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'sick')
-    num_classes = 5
-
-    def __init__(self, dataset_type, cuda):
-        super(SICKDataset, self).__init__(dataset_type, cuda)
-
-    def initialize(self, word_index, embedding):
-        super(SICKDataset, self).initialize(word_index, embedding)
-        new_labels = torch.zeros(self.__len__(), self.num_classes)
-        for i, sim in enumerate(self.labels):
-            ceil, floor = math.ceil(sim), math.floor(sim)
-            if ceil == floor:
-                new_labels[i][floor - 1] = 1
-            else:
-                new_labels[i][floor - 1] = ceil - sim
-                new_labels[i][ceil - 1] = sim - floor
-
-        self.labels = new_labels.cuda() if self.cuda else new_labels
-
-
-class MSRVIDDataset(MPCNNDataset):
-
-    dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'msrvid')
-    num_classes = 6
-
-    def __init__(self, dataset_type, cuda):
-        super(MSRVIDDataset, self).__init__(dataset_type, cuda)
-
-    def initialize(self, word_index, embedding):
-        super(MSRVIDDataset, self).initialize(word_index, embedding)
-        new_labels = torch.zeros(self.__len__(), self.num_classes)
-        for i, sim in enumerate(self.labels):
-            ceil, floor = math.ceil(sim), math.floor(sim)
-            if ceil == floor:
-                new_labels[i][floor] = 1
-            else:
-                new_labels[i][floor] = ceil - sim
-                new_labels[i][ceil] = sim - floor
-
-        self.labels = new_labels.cuda() if self.cuda else new_labels

mp_cnn/evaluation.py

@@ -9,16 +9,16 @@ class MPCNNEvaluatorFactory(object):
     Get the corresponding Evaluator class for a particular dataset.
     """
     @staticmethod
-    def get_evaluator(dataset_name, model, data_loader, batch_size, cuda):
+    def get_evaluator(dataset_cls, model, data_loader, batch_size, device):
         if data_loader is None:
             return None
 
-        if dataset_name == 'sick':
-            return SICKEvaluator(model, data_loader, batch_size, cuda)
-        elif dataset_name == 'msrvid':
-            return MSRVIDEvaluator(model, data_loader, batch_size, cuda)
+        if hasattr(dataset_cls, 'NAME') and dataset_cls.NAME == 'sick':
+            return SICKEvaluator(dataset_cls, model, data_loader, batch_size, device)
+        elif hasattr(dataset_cls, 'NAME') and dataset_cls.NAME == 'msrvid':
+            return MSRVIDEvaluator(dataset_cls, model, data_loader, batch_size, device)
         else:
-            raise ValueError('{} is not a valid dataset.'.format(dataset_name))
+            raise ValueError('{} is not a valid dataset.'.format(dataset_cls))
 
 
 class Evaluator(object):
@@ -26,11 +26,12 @@ class Evaluator(object):
     Evaluates performance of model on a Dataset, using metrics specific to the Dataset.
     """
 
-    def __init__(self, model, data_loader, batch_size, cuda):
+    def __init__(self, dataset_cls, model, data_loader, batch_size, device):
+        self.dataset_cls = dataset_cls
         self.model = model
         self.data_loader = data_loader
         self.batch_size = batch_size
-        self.cuda = cuda
+        self.device = device
 
     def get_scores(self):
         """
@@ -43,74 +44,74 @@ class Evaluator(object):
 
 class SICKEvaluator(Evaluator):
 
-    def __init__(self, model, data_loader, batch_size, cuda):
-        super(SICKEvaluator, self).__init__(model, data_loader, batch_size, cuda)
+    def __init__(self, dataset_cls, model, data_loader, batch_size, device):
+        super(SICKEvaluator, self).__init__(dataset_cls, model, data_loader, batch_size, device)
 
     def get_scores(self):
         self.model.eval()
-        num_classes = self.data_loader.dataset.num_classes
+        num_classes = self.dataset_cls.NUM_CLASSES
         predict_classes = torch.arange(1, num_classes + 1).expand(self.batch_size, num_classes)
-        if self.cuda:
-            predict_classes = predict_classes.cuda()
         test_kl_div_loss = 0
         predictions = []
         true_labels = []
-        for sentences, labels in self.data_loader:
-            sent_a, sent_b = Variable(sentences['a'], volatile=True), Variable(sentences['b'], volatile=True)
-            ext_feats = Variable(sentences['ext_feats'], volatile=True)
-            labels = Variable(labels, volatile=True)
-            output = self.model(sent_a, sent_b, ext_feats)
-            test_kl_div_loss += F.kl_div(output, labels, size_average=False).data[0]
+
+        for batch in self.data_loader:
+            output = self.model(batch.a, batch.b, batch.ext_feats)
+            test_kl_div_loss += F.kl_div(output, batch.label, size_average=False).data[0]
             # handle last batch which might have smaller size
-            if len(predict_classes) != len(sent_a):
-                predict_classes = torch.arange(1, num_classes + 1).expand(len(sent_a), num_classes)
-                if self.cuda:
+            if len(predict_classes) != len(batch.a):
+                predict_classes = torch.arange(1, num_classes + 1).expand(len(batch.a), num_classes)
+
+            if self.data_loader.device != -1:
+                with torch.cuda.device(self.device):
                     predict_classes = predict_classes.cuda()
-            true_labels.append((predict_classes * labels.data).sum(dim=1))
+
+            true_labels.append((predict_classes * batch.label.data).sum(dim=1))
             predictions.append((predict_classes * output.data.exp()).sum(dim=1))
 
             del output
 
         predictions = torch.cat(predictions).cpu().numpy()
         true_labels = torch.cat(true_labels).cpu().numpy()
-        test_kl_div_loss /= len(self.data_loader.dataset)
+        test_kl_div_loss /= len(batch.dataset.examples)
         pearson_r = pearsonr(predictions, true_labels)[0]
         spearman_r = spearmanr(predictions, true_labels)[0]
+
         return [pearson_r, spearman_r, test_kl_div_loss], ['pearson_r', 'spearman_r', 'KL-divergence loss']
 
 
 class MSRVIDEvaluator(Evaluator):
 
-    def __init__(self, model, data_loader, batch_size, cuda):
-        super(MSRVIDEvaluator, self).__init__(model, data_loader, batch_size, cuda)
+    def __init__(self, dataset_cls, model, data_loader, batch_size, device):
+        super(MSRVIDEvaluator, self).__init__(dataset_cls, model, data_loader, batch_size, device)
 
     def get_scores(self):
         self.model.eval()
-        num_classes = self.data_loader.dataset.num_classes
+        num_classes = self.dataset_cls.NUM_CLASSES
         predict_classes = torch.arange(0, num_classes).expand(self.batch_size, num_classes)
-        if self.cuda:
-            predict_classes = predict_classes.cuda()
         test_kl_div_loss = 0
         predictions = []
         true_labels = []
-        for sentences, labels in self.data_loader:
-            sent_a, sent_b = Variable(sentences['a'], volatile=True), Variable(sentences['b'], volatile=True)
-            ext_feats = Variable(sentences['ext_feats'], volatile=True)
-            labels = Variable(labels, volatile=True)
-            output = self.model(sent_a, sent_b, ext_feats)
-            test_kl_div_loss += F.kl_div(output, labels, size_average=False).data[0]
+
+        for batch in self.data_loader:
+            output = self.model(batch.a, batch.b, batch.ext_feats)
+            test_kl_div_loss += F.kl_div(output, batch.label, size_average=False).data[0]
             # handle last batch which might have smaller size
-            if len(predict_classes) != len(sent_a):
-                predict_classes = torch.arange(0, num_classes).expand(len(sent_a), num_classes)
-                if self.cuda:
+            if len(predict_classes) != len(batch.a):
+                predict_classes = torch.arange(0, num_classes).expand(len(batch.a), num_classes)
+
+            if self.data_loader.device != -1:
+                with torch.cuda.device(self.device):
                     predict_classes = predict_classes.cuda()
-            true_labels.append((predict_classes * labels.data).sum(dim=1))
+
+            true_labels.append((predict_classes * batch.label.data).sum(dim=1))
             predictions.append((predict_classes * output.data.exp()).sum(dim=1))
 
             del output
 
         predictions = torch.cat(predictions).cpu().numpy()
         true_labels = torch.cat(true_labels).cpu().numpy()
-        test_kl_div_loss /= len(self.data_loader.dataset)
+        test_kl_div_loss /= len(batch.dataset.examples)
         pearson_r = pearsonr(predictions, true_labels)[0]
+
         return [pearson_r, test_kl_div_loss], ['pearson_r', 'KL-divergence loss']

mp_cnn/main.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import random
 
 import numpy as np
 import torch
@@ -26,9 +27,10 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='PyTorch implementation of Multi-Perspective CNN')
     parser.add_argument('model_outfile', help='file to save final model')
     parser.add_argument('--dataset', help='dataset to use, one of [sick, msrvid]', default='sick')
-    parser.add_argument('--word-vectors-file', help='word vectors file', default=os.path.join(os.pardir, os.pardir, 'data', 'GloVe', 'glove.840B.300d.txt'))
+    parser.add_argument('--word-vectors-dir', help='word vectors directory', default=os.path.join(os.pardir, os.pardir, 'data', 'GloVe'))
+    parser.add_argument('--word-vectors-file', help='word vectors filename', default='glove.840B.300d.txt')
     parser.add_argument('--skip-training', help='will load pre-trained model', action='store_true')
-    parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training (default: false)')
+    parser.add_argument('--device', type=int, default=0, help='GPU device, -1 for CPU (default: 0)')
     parser.add_argument('--sparse-features', action='store_true', default=False, help='use sparse features (default: false)')
     parser.add_argument('--batch-size', type=int, default=64, help='input batch size for training (default: 64)')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train (default: 10)')
@@ -39,7 +41,6 @@ if __name__ == '__main__':
     parser.add_argument('--momentum', type=float, default=0, help='momentum (default: 0)')
     parser.add_argument('--epsilon', type=float, default=1e-8, help='Adam epsilon (default: 1e-8)')
     parser.add_argument('--log-interval', type=int, default=10, help='how many batches to wait before logging training status (default: 10)')
-    parser.add_argument('--sample', type=int, default=0, help='how many examples to take from each dataset, meant for quickly testing entire end-to-end pipeline (default: all)')
     parser.add_argument('--regularization', type=float, default=0.0001, help='Regularization for the optimizer (default: 0.0001)')
     parser.add_argument('--max-window-size', type=int, default=3, help='windows sizes will be [1,max_window_size] and infinity (default: 300)')
     parser.add_argument('--holistic-filters', type=int, default=300, help='number of holistic filters (default: 300)')
@@ -50,20 +51,24 @@ if __name__ == '__main__':
     parser.add_argument('--tensorboard', action='store_true', default=False, help='use TensorBoard to visualize training (default: false)')
     parser.add_argument('--run-label', type=str, help='label to describe run')
     args = parser.parse_args()
-    args.cuda = not args.no_cuda and torch.cuda.is_available()
 
+    random.seed(args.seed)
     np.random.seed(args.seed)
     torch.manual_seed(args.seed)
-    if args.cuda:
+    if args.device != -1:
         torch.cuda.manual_seed(args.seed)
 
-    train_loader, test_loader, dev_loader = MPCNNDatasetFactory.get_dataset(args.dataset, args.word_vectors_file, args.batch_size, args.cuda, args.sample)
+    dataset_cls, embedding, train_loader, test_loader, dev_loader \
+        = MPCNNDatasetFactory.get_dataset(args.dataset, args.word_vectors_dir, args.word_vectors_file, args.batch_size, args.device)
 
     filter_widths = list(range(1, args.max_window_size + 1)) + [np.inf]
-    input_channels = 300
-    model = MPCNN(input_channels, args.holistic_filters, args.per_dim_filters, filter_widths, args.hidden_units, train_loader.dataset.num_classes, args.dropout, args.sparse_features)
-    if args.cuda:
-        model.cuda()
+    model = MPCNN(embedding, args.holistic_filters, args.per_dim_filters, filter_widths,
+                    args.hidden_units, dataset_cls.NUM_CLASSES, args.dropout, args.sparse_features)
+
+    if args.device != -1:
+        with torch.cuda.device(args.device):
+            model.cuda()
+
     optimizer = None
     if args.optimizer == 'adam':
         optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.regularization, eps=args.epsilon)
@@ -71,14 +76,14 @@ if __name__ == '__main__':
         optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.regularization)
     else:
         raise ValueError('optimizer not recognized: it should be either adam or sgd')
-    train_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, train_loader, args.batch_size, args.cuda)
-    test_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, test_loader, args.batch_size, args.cuda)
-    dev_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, dev_loader, args.batch_size, args.cuda)
+
+    train_evaluator = MPCNNEvaluatorFactory.get_evaluator(dataset_cls, model, train_loader, args.batch_size, args.device)
+    test_evaluator = MPCNNEvaluatorFactory.get_evaluator(dataset_cls, model, test_loader, args.batch_size, args.device)
+    dev_evaluator = MPCNNEvaluatorFactory.get_evaluator(dataset_cls, model, dev_loader, args.batch_size, args.device)
 
     trainer_config = {
         'optimizer': optimizer,
         'batch_size': args.batch_size,
-        'sample': args.sample,
         'log_interval': args.log_interval,
         'model_outfile': args.model_outfile,
         'lr_reduce_factor': args.lr_reduce_factor,
@@ -97,8 +102,8 @@ if __name__ == '__main__':
         trainer.train(args.epochs)
 
     model = torch.load(args.model_outfile)
-    test_evaluator = MPCNNEvaluatorFactory.get_evaluator(args.dataset, model, test_loader, args.batch_size, args.cuda)
-    scores, metric_names = test_evaluator.get_scores()
+    saved_model_evaluator = MPCNNEvaluatorFactory.get_evaluator(dataset_cls, model, test_loader, args.batch_size, args.device)
+    scores, metric_names = saved_model_evaluator.get_scores()
     logger.info('Evaluation metrics for test')
     logger.info('\t'.join([' '] + metric_names))
     logger.info('\t'.join(['test'] + list(map(str, scores))))

mp_cnn/model.py

@@ -6,10 +6,10 @@ import torch.nn.functional as F
 
 class MPCNN(nn.Module):
 
-    def __init__(self, n_word_dim, n_holistic_filters, n_per_dim_filters, filter_widths, hidden_layer_units, num_classes, dropout, ext_feats):
+    def __init__(self, embedding, n_holistic_filters, n_per_dim_filters, filter_widths, hidden_layer_units, num_classes, dropout, ext_feats):
         super(MPCNN, self).__init__()
-
-        self.n_word_dim = n_word_dim
+        self.embedding = embedding
+        self.n_word_dim = embedding.weight.size(1)
         self.n_holistic_filters = n_holistic_filters
         self.n_per_dim_filters = n_per_dim_filters
         self.filter_widths = filter_widths
@@ -22,12 +22,12 @@ class MPCNN(nn.Module):
                 continue
 
             holistic_conv_layers.append(nn.Sequential(
-                nn.Conv1d(n_word_dim, n_holistic_filters, ws),
+                nn.Conv1d(self.n_word_dim, n_holistic_filters, ws),
                 nn.Tanh()
             ))
 
             per_dim_conv_layers.append(nn.Sequential(
-                nn.Conv1d(n_word_dim, n_word_dim * n_per_dim_filters, ws, groups=n_word_dim),
+                nn.Conv1d(self.n_word_dim, self.n_word_dim * n_per_dim_filters, ws, groups=self.n_word_dim),
                 nn.Tanh()
             ))
 
@@ -35,7 +35,7 @@ class MPCNN(nn.Module):
         self.per_dim_conv_layers = nn.ModuleList(per_dim_conv_layers)
 
         # compute number of inputs to first hidden layer
-        COMP_1_COMPONENTS_HOLISTIC, COMP_1_COMPONENTS_PER_DIM, COMP_2_COMPONENTS = 2 + n_holistic_filters, 2 + n_word_dim, 2
+        COMP_1_COMPONENTS_HOLISTIC, COMP_1_COMPONENTS_PER_DIM, COMP_2_COMPONENTS = 2 + n_holistic_filters, 2 + self.n_word_dim, 2
         EXT_FEATS = 4 if ext_feats else 0
         n_feat_h = 3 * len(self.filter_widths) * COMP_2_COMPONENTS
         n_feat_v = (
@@ -61,24 +61,25 @@ class MPCNN(nn.Module):
         block_b = {}
         for ws in self.filter_widths:
             if np.isinf(ws):
+                sent_flattened, sent_flattened_size = sent.contiguous().view(sent.size(0), 1, -1), sent.size(1) * sent.size(2)
                 block_a[ws] = {
-                    'max': F.max_pool1d(sent.view(sent.size(0), 1, -1), sent.size(1) * sent.size(2)).view(sent.size(0), -1),
-                    'min': F.max_pool1d(-1 * sent.view(sent.size(0), 1, -1), sent.size(1) * sent.size(2)).view(sent.size(0), -1),
-                    'mean': F.avg_pool1d(sent.view(sent.size(0), 1, -1), sent.size(1) * sent.size(2)).view(sent.size(0), -1)
+                    'max': F.max_pool1d(sent_flattened, sent_flattened_size).view(sent.size(0), -1),
+                    'min': F.max_pool1d(-1 * sent_flattened, sent_flattened_size).view(sent.size(0), -1),
+                    'mean': F.avg_pool1d(sent_flattened, sent_flattened_size).view(sent.size(0), -1)
                 }
                 continue
 
             holistic_conv_out = self.holistic_conv_layers[ws - 1](sent)
             block_a[ws] = {
-                'max': F.max_pool1d(holistic_conv_out, holistic_conv_out.size(2)).view(-1, self.n_holistic_filters),
-                'min': F.max_pool1d(-1 * holistic_conv_out, holistic_conv_out.size(2)).view(-1, self.n_holistic_filters),
-                'mean': F.avg_pool1d(holistic_conv_out, holistic_conv_out.size(2)).view(-1, self.n_holistic_filters)
+                'max': F.max_pool1d(holistic_conv_out, holistic_conv_out.size(2)).contiguous().view(-1, self.n_holistic_filters),
+                'min': F.max_pool1d(-1 * holistic_conv_out, holistic_conv_out.size(2)).contiguous().view(-1, self.n_holistic_filters),
+                'mean': F.avg_pool1d(holistic_conv_out, holistic_conv_out.size(2)).contiguous().view(-1, self.n_holistic_filters)
             }
 
             per_dim_conv_out = self.per_dim_conv_layers[ws - 1](sent)
             block_b[ws] = {
-                'max': F.max_pool1d(per_dim_conv_out, per_dim_conv_out.size(2)).view(-1, self.n_word_dim, self.n_per_dim_filters),
-                'min': F.max_pool1d(-1 * per_dim_conv_out, per_dim_conv_out.size(2)).view(-1, self.n_word_dim, self.n_per_dim_filters)
+                'max': F.max_pool1d(per_dim_conv_out, per_dim_conv_out.size(2)).contiguous().view(-1, self.n_word_dim, self.n_per_dim_filters),
+                'min': F.max_pool1d(-1 * per_dim_conv_out, per_dim_conv_out.size(2)).contiguous().view(-1, self.n_word_dim, self.n_per_dim_filters)
             }
         return block_a, block_b
 
@@ -89,7 +90,7 @@ class MPCNN(nn.Module):
                 x1 = sent1_block_a[ws][pool]
                 x2 = sent2_block_a[ws][pool]
                 batch_size = x1.size()[0]
-                comparison_feats.append(F.cosine_similarity(x1, x2).view(batch_size, 1))
+                comparison_feats.append(F.cosine_similarity(x1, x2).contiguous().view(batch_size, 1))
                 comparison_feats.append(F.pairwise_distance(x1, x2))
         return torch.cat(comparison_feats, dim=1)
 
@@ -103,7 +104,7 @@ class MPCNN(nn.Module):
                 for ws2 in self.filter_widths:
                     x2 = sent2_block_a[ws2][pool]
                     if (not np.isinf(ws1) and not np.isinf(ws2)) or (np.isinf(ws1) and np.isinf(ws2)):
-                        comparison_feats.append(F.cosine_similarity(x1, x2).view(batch_size, 1))
+                        comparison_feats.append(F.cosine_similarity(x1, x2).contiguous().view(batch_size, 1))
                         comparison_feats.append(F.pairwise_distance(x1, x2))
                         comparison_feats.append(torch.abs(x1 - x2))
 
@@ -115,13 +116,17 @@ class MPCNN(nn.Module):
                     x1 = oG_1B[:, :, i]
                     x2 = oG_2B[:, :, i]
                     batch_size = x1.size()[0]
-                    comparison_feats.append(F.cosine_similarity(x1, x2).view(batch_size, 1))
+                    comparison_feats.append(F.cosine_similarity(x1, x2).contiguous().view(batch_size, 1))
                     comparison_feats.append(F.pairwise_distance(x1, x2))
                     comparison_feats.append(torch.abs(x1 - x2))
 
         return torch.cat(comparison_feats, dim=1)
 
-    def forward(self, sent1, sent2, ext_feats):
+    def forward(self, sent1_idx, sent2_idx, ext_feats=None):
+        # Select embedding
+        sent1 = self.embedding(sent1_idx).transpose(1, 2)
+        sent2 = self.embedding(sent2_idx).transpose(1, 2)
+
         # Sentence modeling module
         sent1_block_a, sent1_block_b = self._get_blocks_for_sentence(sent1)
         sent2_block_a, sent2_block_b = self._get_blocks_for_sentence(sent2)

mp_cnn/preprocessing.py

@@ -1,55 +0,0 @@
-"""
-Preprocessing utilities such as preparing embeddings.
-"""
-import os
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-
-def get_vocab(data_dir):
-    """
-    Get vocabulary as set of words.
-    """
-    vocab = set()
-    with open(os.path.join(data_dir, 'vocab-cased.txt'), 'r') as f:
-        for line in f:
-            word = line.rstrip()
-            vocab.add(word)
-    return vocab
-
-
-def get_embedding_index(vocab, glove_file):
-    """
-    Get dictionary mapping word to its word vector.
-    """
-    embedding_index = {}
-    with open(glove_file, 'r') as f:
-        for line in f:
-            parts = line.split()
-            word = parts[0]
-            if word not in vocab or len(parts) != 301:
-                continue
-            vec = np.asarray(parts[1:], dtype='float32')
-            embedding_index[word] = vec
-    return embedding_index
-
-
-def get_glove_embedding(glove_file, data_dir):
-    """
-    Get embedding for the words in the data set.
-    """
-    vocab = get_vocab(data_dir)
-    embedding_index = get_embedding_index(vocab, glove_file)
-    word_index = {w: i for i, w in enumerate(embedding_index.keys())}
-
-    embedding_matrix = np.zeros((len(word_index), 300))
-    for word, i in word_index.items():
-        embedding_matrix[i] = embedding_index.get(word)
-
-    embedding_tensor = torch.from_numpy(embedding_matrix)
-    embedding = nn.Embedding(len(embedding_index), 300)
-    embedding.weight = nn.Parameter(embedding_tensor)
-    return word_index, embedding
-

mp_cnn/train.py

@@ -44,7 +44,6 @@ class Trainer(object):
         self.optimizer = trainer_config['optimizer']
         self.train_loader = train_loader
         self.batch_size = trainer_config['batch_size']
-        self.sample = trainer_config['sample']
         self.log_interval = trainer_config['log_interval']
         self.model_outfile = trainer_config['model_outfile']
         self.lr_reduce_factor = trainer_config['lr_reduce_factor']
@@ -80,21 +79,18 @@ class SICKTrainer(Trainer):
     def train_epoch(self, epoch):
         self.model.train()
         total_loss = 0
-        for batch_idx, (sentences, labels) in enumerate(self.train_loader):
-            sent_a, sent_b = Variable(sentences['a']), Variable(sentences['b'])
-            ext_feats = Variable(sentences['ext_feats'])
-            labels = Variable(labels)
+        for batch_idx, batch in enumerate(self.train_loader):
             self.optimizer.zero_grad()
-            output = self.model(sent_a, sent_b, ext_feats)
-            loss = F.kl_div(output, labels)
+            output = self.model(batch.a, batch.b, batch.ext_feats)
+            loss = F.kl_div(output, batch.label)
             total_loss += loss.data[0]
             loss.backward()
             self.optimizer.step()
             if batch_idx % self.log_interval == 0:
                 logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
-                    epoch, min(batch_idx * self.batch_size, len(self.train_loader.dataset)),
-                    len(self.train_loader.dataset) if not self.sample else self.sample,
-                    100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0])
+                    epoch, min(batch_idx * self.batch_size, len(batch.dataset.examples)),
+                    len(batch.dataset.examples),
+                    100. * batch_idx / (len(self.train_loader)), loss.data[0])
                 )
 
         if self.use_tensorboard:
@@ -146,34 +142,33 @@ class MSRVIDTrainer(Trainer):
 
     def train_epoch(self, epoch):
         self.model.train()
+        total_loss = 0
 
         # since MSRVID doesn't have validation set, we manually leave-out some training data for validation
-        batches = math.ceil(len(self.train_loader.dataset) / self.batch_size)
+        batches = math.ceil(len(self.train_loader.dataset.examples) / self.batch_size)
         start_val_batch = math.floor(0.8 * batches)
         left_out_val_a, left_out_val_b = [], []
-        left_out_ext_feats = []
+        left_out_val_ext_feats = []
         left_out_val_labels = []
-        total_loss = 0
-        for batch_idx, (sentences, labels) in enumerate(self.train_loader):
-            sent_a, sent_b = Variable(sentences['a']), Variable(sentences['b'])
-            ext_feats = Variable(sentences['ext_feats'])
-            labels = Variable(labels)
+
+        for batch_idx, batch in enumerate(self.train_loader):
             if batch_idx >= start_val_batch:
-                left_out_val_a.append(sent_a)
-                left_out_val_b.append(sent_b)
-                left_out_val_labels.append(labels)
+                left_out_val_a.append(batch.a)
+                left_out_val_b.append(batch.b)
+                left_out_val_ext_feats.append(batch.ext_feats)
+                left_out_val_labels.append(batch.label)
                 continue
             self.optimizer.zero_grad()
-            output = self.model(sent_a, sent_b, ext_feats)
-            loss = F.kl_div(output, labels)
+            output = self.model(batch.a, batch.b, batch.ext_feats)
+            loss = F.kl_div(output, batch.label)
             total_loss += loss.data[0]
             loss.backward()
             self.optimizer.step()
             if batch_idx % self.log_interval == 0:
                 logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
-                    epoch, min(batch_idx * self.batch_size, len(self.train_loader.dataset)),
-                    len(self.train_loader.dataset) if not self.sample else self.sample,
-                    100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0])
+                    epoch, min(batch_idx * self.batch_size, len(batch.dataset.examples)),
+                    len(batch.dataset.examples),
+                    100. * batch_idx / (len(self.train_loader)), loss.data[0])
                 )
 
         self.evaluate(self.train_evaluator, 'train')
@@ -181,7 +176,7 @@ class MSRVIDTrainer(Trainer):
         if self.use_tensorboard:
             self.writer.add_scalar('msrvid/train/kl_div_loss', total_loss, epoch)
 
-        return left_out_val_a, left_out_val_b, left_out_ext_feats, left_out_val_labels
+        return left_out_val_a, left_out_val_b, left_out_val_ext_feats, left_out_val_labels
 
     def train(self, epochs):
         scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=self.lr_reduce_factor, patience=self.patience)
@@ -194,18 +189,25 @@ class MSRVIDTrainer(Trainer):
             left_out_a, left_out_b, left_out_ext_feats, left_out_label = self.train_epoch(epoch)
 
             # manually evaluating the validating set
-            left_out_a = torch.cat(left_out_a)
-            left_out_b = torch.cat(left_out_b)
-            left_out_ext_feats = torch.cat(left_out_ext_feats)
-            left_out_label = torch.cat(left_out_label)
-            output = self.model(left_out_a, left_out_b, left_out_ext_feats)
-            val_kl_div_loss = F.kl_div(output, left_out_label).data[0]
-            predict_classes = torch.arange(0, 6).expand(len(left_out_a), 6).cuda()
-            true_labels = (predict_classes * left_out_label.data).sum(dim=1)
-            predictions = (predict_classes * output.data.exp()).sum(dim=1)
-            predictions = predictions.cpu().numpy()
-            true_labels = true_labels.cpu().numpy()
+            all_predictions, all_true_labels = [], []
+            val_kl_div_loss = 0
+            for i in range(len(left_out_a)):
+                output = self.model(left_out_a[i], left_out_b[i], left_out_ext_feats[i])
+                val_kl_div_loss += F.kl_div(output, left_out_label[i], size_average=False).data[0]
+                predict_classes = torch.arange(0, self.train_loader.dataset.NUM_CLASSES).expand(len(left_out_a[i]), self.train_loader.dataset.NUM_CLASSES)
+                if self.train_loader.device != -1:
+                    with torch.cuda.device(self.train_loader.device):
+                        predict_classes = predict_classes.cuda()
+
+                predictions = (predict_classes * output.data.exp()).sum(dim=1)
+                true_labels = (predict_classes * left_out_label[i].data).sum(dim=1)
+                all_predictions.append(predictions)
+                all_true_labels.append(true_labels)
+
+            predictions = torch.cat(all_predictions).cpu().numpy()
+            true_labels = torch.cat(all_true_labels).cpu().numpy()
             pearson_r = pearsonr(predictions, true_labels)[0]
+            val_kl_div_loss /= len(predictions)
 
             if self.use_tensorboard:
                 self.writer.add_scalar('msrvid/dev/pearson_r', pearson_r, epoch)