diff --git a/vdpwi/__main__.py b/vdpwi/__main__.py new file mode 100644 index 0000000..05a56be --- /dev/null +++ b/vdpwi/__main__.py @@ -0,0 +1,139 @@ +from collections import namedtuple + +from tqdm import tqdm +import numpy as np +import scipy.stats as stats +import torch +import torch.optim as optim +import torch.nn as nn +import torch.nn.functional as F +import torch.utils as utils + +from utils.log import LogWriter +import data +import model as mod + +Context = namedtuple("Context", "model, train_loader, dev_loader, test_loader, optimizer, criterion, params, log_writer") +EvaluateResult = namedtuple("EvaluateResult", "pearsonr, spearmanr") + +def create_context(config): + def collate_fn(batch): + emb1 = [] + emb2 = [] + labels = [] + cmp_labels = [] + pad_cube = [] + max_len1 = 0; max_len2 = 0 + + for s1, s2, l, cl in batch: + emb1.append(s1) + emb2.append(s2) + max_len1 = max(max_len1, len(s1)) + max_len2 = max(max_len2, len(s2)) + labels.append(l) + cmp_labels.append(cl) + + for s1, s2 in zip(emb1, emb2): + pad1 = (max_len1 - len(s1)) + pad2 = (max_len2 - len(s2)) + pad_mask = np.ones((max_len1, max_len2)) + pad_mask[:len(s1), :len(s2)] = 0 + pad_cube.append(pad_mask) + s1.extend([embedding.weight.size(0) - 1] * pad1) + s2.extend([embedding.weight.size(0) - 1] * pad2) + + pad_cube = np.array(pad_cube) + emb1 = torch.LongTensor(emb1) + emb2 = torch.LongTensor(emb2) + labels = torch.Tensor(labels) + emb1 = torch.autograd.Variable(emb1, requires_grad=False) + emb2 = torch.autograd.Variable(emb2, requires_grad=False) + labels = torch.autograd.Variable(labels, requires_grad=False) + pad_cube = torch.autograd.Variable(torch.from_numpy(pad_cube).float(), requires_grad=False) + if not config.cpu: + emb1 = emb1.cuda() + emb2 = emb2.cuda() + labels = labels.cuda() + pad_cube = pad_cube.cuda() + return emb1, emb2, labels, pad_cube, cmp_labels + + embedding, (train_set, dev_set, test_set) = data.load_dataset(config.dataset) + model = mod.VDPWIModel(embedding, config) + if config.restore: + model.load(config.input_file) + if not config.cpu: + model = model.cuda() + + train_loader = utils.data.DataLoader(train_set, shuffle=True, batch_size=config.mbatch_size, collate_fn=collate_fn) + dev_loader = utils.data.DataLoader(dev_set, batch_size=1, collate_fn=collate_fn) + test_loader = utils.data.DataLoader(test_set, batch_size=1, collate_fn=collate_fn) + + params = list(filter(lambda x: x.requires_grad, model.parameters())) + if config.optimizer == "adam": + optimizer = optim.Adam(params, lr=config.lr, weight_decay=config.weight_decay) + elif config.optimizer == "sgd": + optimizer = optim.SGD(params, lr=config.lr, momentum=config.momentum, weight_decay=config.weight_decay) + elif config.optimizer == "rmsprop": + optimizer = optim.RMSprop(params, lr=config.lr, alpha=config.decay, momentum=config.momentum, weight_decay=config.weight_decay) + criterion = nn.KLDivLoss() + log_writer = LogWriter() + return Context(model, train_loader, dev_loader, test_loader, optimizer, criterion, params, log_writer) + +def test(config): + context = create_context(config) + result = evaluate(context, context.test_loader) + print("Final test result: {}".format(result)) + +def evaluate(context, data_loader): + model = context.model + model.eval() + predictions = [] + true_labels = [] + for sent1, sent2, _, pad_cube, truth in data_loader: + scores = model(sent1, sent2, pad_cube) + scores = F.softmax(scores).cpu().data.numpy()[0] + prediction = np.dot(np.arange(1, len(scores) + 1), scores) + predictions.append(prediction); true_labels.append(truth[0][0]) + + pearsonr = stats.pearsonr(predictions, true_labels)[0] + spearmanr = stats.spearmanr(predictions, true_labels)[0] + context.log_writer.log_dev_metrics(pearsonr, spearmanr) + return EvaluateResult(pearsonr, spearmanr) + +def train(config): + context = create_context(config) + context.log_writer.log_hyperparams() + best_dev_pr = 0 + for epoch_no in range(config.n_epochs): + print("Epoch number: {}".format(epoch_no + 1)) + loader_wrapper = tqdm(context.train_loader, total=len(context.train_loader), desc="Loss") + context.model.train() + loss = 0 + for sent1, sent2, label_pmf, pad_cube, _ in loader_wrapper: + context.optimizer.zero_grad() + scores = F.log_softmax(context.model(sent1, sent2, pad_cube)) + + loss = context.criterion(scores, label_pmf) + loss.backward() + nn.utils.clip_grad_norm(context.params, config.clip_norm) + context.optimizer.step() + + loss = loss.cpu().data[0] + loader_wrapper.set_description("Loss: {:<8}".format(round(loss, 5))) + context.log_writer.log_train_loss(loss) + result = evaluate(context, context.dev_loader) + print("Dev result: {}".format(result)) + if best_dev_pr < result.pearsonr: + best_dev_pr = result.pearsonr + print("Saving best model...") + context.model.save(config.output_file) + +def main(): + config = data.Configs.base_config() + if config.mode == "train": + train(config) + elif config.mode == "test": + test(config) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/vdpwi/data.py b/vdpwi/data.py new file mode 100644 index 0000000..04cd560 --- /dev/null +++ b/vdpwi/data.py @@ -0,0 +1,104 @@ +import argparse +import os + +import torch +import torch.nn as nn +import torch.utils.data as data + +class Configs(object): + @staticmethod + def base_config(): + parser = argparse.ArgumentParser() + parser.add_argument("--classifier", type=str, default="vdpwi", choices=["vdpwi", "resnet"]) + parser.add_argument("--clip_norm", type=float, default=50) + parser.add_argument("--cpu", action="store_true", default=False) + parser.add_argument("--dataset", type=str, default="sick", choices=["sick"]) + parser.add_argument("--decay", type=float, default=0.95) + parser.add_argument("--input_file", type=str, default="local_saves/model.pt") + parser.add_argument("--lr", type=float, default=5E-4) + parser.add_argument("--mbatch_size", type=int, default=16) + parser.add_argument("--mode", type=str, default="train", choices=["train", "test"]) + parser.add_argument("--momentum", type=float, default=0.1) + parser.add_argument("--n_epochs", type=int, default=35) + parser.add_argument("--n_labels", type=int, default=5) + parser.add_argument("--optimizer", type=str, default="rmsprop", choices=["adam", "sgd", "rmsprop"]) + parser.add_argument("--output_file", type=str, default="local_saves/model.pt") + parser.add_argument("--res_fmaps", type=int, default=32) + parser.add_argument("--res_layers", type=int, default=16) + parser.add_argument("--restore", action="store_true", default=False) + parser.add_argument("--rnn_hidden_dim", type=int, default=250) + parser.add_argument("--weight_decay", type=float, default=1E-5) + parser.add_argument("--wordvecs_file", type=str, default="local_data/glove/glove.840B.300d.txt") + return parser.parse_known_args()[0] + + @staticmethod + def sick_config(): + parser = argparse.ArgumentParser() + parser.add_argument("--n_labels", type=int, default=5) + parser.add_argument("--sick_cache", type=str, default="local_data/sick/.vec-cache") + parser.add_argument("--sick_data", type=str, default="local_data/sick") + return parser.parse_known_args()[0] + +class LabeledEmbeddedDataset(data.Dataset): + def __init__(self, sentence_indices1, sentence_indices2, labels, compare_labels=None): + assert len(sentence_indices1) == len(labels) == len(sentence_indices2) + self.sentence_indices1 = sentence_indices1 + self.sentence_indices2 = sentence_indices2 + self.labels = labels + self.compare_labels = compare_labels + + def __getitem__(self, idx): + cmp_lbl = None if self.compare_labels is None else self.compare_labels[idx] + return self.sentence_indices1[idx], self.sentence_indices2[idx], self.labels[idx], cmp_lbl + + def __len__(self): + return len(self.labels) + +def load_sick(): + config = Configs.sick_config() + def fetch_indices(name): + sentence_indices = [] + filename = os.path.join(config.sick_data, dataset, name) + with open(filename) as f: + for line in f: + indices = [embed_ids.get(word, -1) for word in line.strip().split()] + indices = list(filter(lambda x: x >= 0, indices)) + sentence_indices.append(indices) + return sentence_indices + + def read_labels(filename): + labels = [] + with open(filename) as f: + for line in f: + labels.append([float(val) for val in line.split()]) + return labels + + sets = [] + embeddings = [] + embed_ids = {} + with open(os.path.join(config.sick_cache)) as f: + for i, line in enumerate(f): + word, vec = line.split(" ", 1) + vec = list(map(float, vec.strip().split())) + embed_ids[word] = i + embeddings.append(vec) + padding_idx = len(embeddings) + embeddings.append([0.0] * 300) + + for dataset in ("train", "dev", "test"): + sparse_filename = os.path.join(config.sick_data, dataset, "sim_sparse.txt") + truth_filename = os.path.join(config.sick_data, dataset, "sim.txt") + sparse_labels = read_labels(sparse_filename) + cmp_labels = read_labels(truth_filename) + indices1 = fetch_indices("a.toks") + indices2 = fetch_indices("b.toks") + sets.append(LabeledEmbeddedDataset(indices1, indices2, sparse_labels, cmp_labels)) + embedding = nn.Embedding(len(embeddings), 300) + embedding.weight.data.copy_(torch.Tensor(embeddings)) + embedding.weight.requires_grad = False + return embedding, sets + +def load_dataset(dataset): + return _loaders[dataset]() + +_loaders = dict(sick=load_sick) diff --git a/vdpwi/model.py b/vdpwi/model.py new file mode 100644 index 0000000..c31bcbc --- /dev/null +++ b/vdpwi/model.py @@ -0,0 +1,168 @@ +from torch.autograd import Variable +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.models as models +import numpy as np + +class SerializableModule(nn.Module): + def __init__(self): + super().__init__() + + def save(self, filename): + torch.save(self.state_dict(), filename) + + def load(self, filename): + self.load_state_dict(torch.load(filename, map_location=lambda storage, loc: storage)) + +def hard_pad2d(x, pad): + def pad_side(idx): + pad_len = max(pad - x.size(idx), 0) + return [0, pad_len] + padding = pad_side(3) + padding.extend(pad_side(2)) + x = F.pad(x, padding) + return x[:, :, :pad, :pad] + +class ResNet(SerializableModule): + def __init__(self, config): + super().__init__() + n_layers = config.res_layers + n_maps = config.res_fmaps + n_labels = config.n_labels + self.conv0 = nn.Conv2d(12, n_maps, (3, 3), padding=1) + self.convs = [nn.Conv2d(n_maps, n_maps, (3, 3), padding=1) for _ in range(n_layers)] + self.output = nn.Linear(n_maps, n_labels) + self.input_len = None + for i, conv in enumerate(self.convs): + self.add_module("conv{}".format(i + 1), conv) + + def forward(self, x): + x = F.relu(self.conv0(x)) + old_x = x + for i, conv in enumerate(self.convs): + x = F.relu(conv(x)) + if i % 2 == 1: + x += old_x + old_x = x + x = torch.mean(x.view(x.size(0), x.size(1), -1), 2) + return self.output(x) + +class VDPWIConvNet(SerializableModule): + def __init__(self, config): + super().__init__() + def make_conv(n_in, n_out): + conv = nn.Conv2d(n_in, n_out, 3, padding=1) + conv.bias.data.zero_() + nn.init.xavier_normal(conv.weight) + return conv + self.conv1 = make_conv(12, 128) + self.conv2 = make_conv(128, 164) + self.conv3 = make_conv(164, 192) + self.conv4 = make_conv(192, 192) + self.conv5 = make_conv(192, 128) + self.maxpool2 = nn.MaxPool2d(2, ceil_mode=True) + self.dnn = nn.Linear(128, 128) + self.output = nn.Linear(128, config.n_labels) + self.input_len = 32 + + def forward(self, x): + x = hard_pad2d(x, self.input_len) + pool_final = nn.MaxPool2d(2, ceil_mode=True) if x.size(2) == 32 else nn.MaxPool2d(3, 1, ceil_mode=True) + x = self.maxpool2(F.relu(self.conv1(x))) + x = self.maxpool2(F.relu(self.conv2(x))) + x = self.maxpool2(F.relu(self.conv3(x))) + x = self.maxpool2(F.relu(self.conv4(x))) + x = pool_final(F.relu(self.conv5(x))) + x = F.relu(self.dnn(x.view(x.size(0), -1))) + return self.output(x) + +class VDPWIModel(SerializableModule): + def __init__(self, embedding, config): + super().__init__() + self.hidden_dim = config.rnn_hidden_dim + self.rnn = nn.LSTM(300, self.hidden_dim, 1, batch_first=True) + self.embedding = embedding + self.use_cuda = not config.cpu + if config.classifier == "vdpwi": + self.classifier_net = VDPWIConvNet(config) + elif config.classifier == "resnet": + self.classifier_net = ResNet(config) + + def compute_sim_cube(self, seq1, seq2): + def compute_sim(prism1, prism2): + prism1_len = prism1.norm(dim=3) + prism2_len = prism2.norm(dim=3) + + dot_prod = torch.matmul(prism1.unsqueeze(3), prism2.unsqueeze(4)) + dot_prod = dot_prod.squeeze(3).squeeze(3) + cos_dist = dot_prod / (prism1_len * prism2_len + 1E-8) + l2_dist = -((prism1 - prism2).norm(dim=3)) + return torch.stack([dot_prod, cos_dist, l2_dist], 1) + + def compute_prism(seq1, seq2): + prism1 = seq1.repeat(seq2.size(1), 1, 1, 1) + prism2 = seq2.repeat(seq1.size(1), 1, 1, 1) + prism1 = prism1.permute(1, 2, 0, 3).contiguous() + prism2 = prism2.permute(1, 0, 2, 3).contiguous() + return compute_sim(prism1, prism2) + + sim_cube = Variable(torch.Tensor(seq1.size(0), 12, seq1.size(1), seq2.size(1))) + if self.use_cuda: + sim_cube = sim_cube.cuda() + seq1_f = seq1[:, :, :self.hidden_dim] + seq1_b = seq1[:, :, self.hidden_dim:] + seq2_f = seq2[:, :, :self.hidden_dim] + seq2_b = seq2[:, :, self.hidden_dim:] + sim_cube[:, 0:3] = compute_prism(seq1, seq2) + sim_cube[:, 3:6] = compute_prism(seq1_f, seq2_f) + sim_cube[:, 6:9] = compute_prism(seq1_b, seq2_b) + sim_cube[:, 9:12] = compute_prism(seq1_f + seq1_b, seq2_f + seq2_b) + return sim_cube + + def compute_focus_cube(self, sim_cube, pad_cube): + neg_magic = -10000 + pad_cube = pad_cube.repeat(12, 1, 1, 1) + pad_cube = pad_cube.permute(1, 0, 2, 3).contiguous() + sim_cube = neg_magic * pad_cube + sim_cube + mask = Variable(torch.Tensor(*sim_cube.size())) + if self.use_cuda: + mask = mask.cuda() + mask[:, :, :, :] = 0.1 + + def build_mask(index): + max_mask = sim_cube[:, index].clone() + for _ in range(min(sim_cube.size(2), sim_cube.size(3))): + values, indices = torch.max(max_mask.view(sim_cube.size(0), -1), 1) + row_indices = indices / sim_cube.size(3) + col_indices = indices % sim_cube.size(3) + row_indices = row_indices.unsqueeze(1) + col_indices = col_indices.unsqueeze(1).unsqueeze(1) + for i, (row_i, col_i, val) in enumerate(zip(row_indices, col_indices, values)): + if val < neg_magic / 2: + continue + mask[i, :, row_i, col_i] = 1 + max_mask[i, row_i, :] = neg_magic + max_mask[i, :, col_i] = neg_magic + build_mask(9) + build_mask(10) + focus_cube = mask * sim_cube * (1 - pad_cube) + return focus_cube + + def forward(self, x1, x2, pad_cube): + x1 = self.embedding(x1) + x2 = self.embedding(x2) + seq1f, _ = self.rnn(x1) + seq2f, _ = self.rnn(x2) + seq1b, _ = self.rnn(torch.cat(x1.split(1, 1)[::-1], 1)) + seq2b, _ = self.rnn(torch.cat(x2.split(1, 1)[::-1], 1)) + seq1 = torch.cat([seq1f, seq1b], 2) + seq2 = torch.cat([seq2f, seq2b], 2) + sim_cube = self.compute_sim_cube(seq1, seq2) + truncate = self.classifier_net.input_len + if truncate is not None: + sim_cube = sim_cube[:, :, :truncate, :truncate].contiguous() + pad_cube = pad_cube[:, :truncate, :truncate].contiguous() + focus_cube = self.compute_focus_cube(sim_cube, pad_cube) + logits = self.classifier_net(focus_cube) + return logits diff --git a/vdpwi/train_default.sh b/vdpwi/train_default.sh new file mode 100755 index 0000000..230a237 --- /dev/null +++ b/vdpwi/train_default.sh @@ -0,0 +1,2 @@ +#!/bin/sh +python . --clip_norm 50 --decay 0.95 --lr 1E-4 --mbatch_size 1 --momentum 0 --optimizer rmsprop --weight_decay 0 diff --git a/vdpwi/utils/__init__.py b/vdpwi/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/vdpwi/utils/log.py b/vdpwi/utils/log.py new file mode 100644 index 0000000..82f5eef --- /dev/null +++ b/vdpwi/utils/log.py @@ -0,0 +1,26 @@ +import datetime +import sys + +from tensorboardX import SummaryWriter + +class LogWriter(object): + def __init__(self, run_name_fmt="run_{}"): + self.writer = SummaryWriter() + self.run_name = run_name_fmt.format(datetime.datetime.now().strftime("%Y-%m-%d-%H:%M:%S")) + self.train_idx = 0 + self.dev_idx = 0 + + def log_hyperparams(self): + self.writer.add_text("{}/hyperparams".format(self.run_name), " ".join(sys.argv)) + + def log_train_loss(self, loss): + self.writer.add_scalar("{}/train_loss".format(self.run_name), loss, self.train_idx) + self.train_idx += 1 + + def log_dev_metrics(self, pearsonr, spearmanr): + results = dict(pearsonr=pearsonr, spearmanr=spearmanr) + self.writer.add_scalars("{}/dev_metrics".format(self.run_name), results, self.dev_idx) + self.dev_idx += 1 + + def next(self): + self.i += 1 diff --git a/vdpwi/utils/preprocess.py b/vdpwi/utils/preprocess.py new file mode 100644 index 0000000..c0b4e77 --- /dev/null +++ b/vdpwi/utils/preprocess.py @@ -0,0 +1,78 @@ +import argparse +import os + +from scipy.special import erf +from scipy.stats import truncnorm +import numpy as np + +import data + +def build_vector_cache(glove_filename, vec_cache_filename, vocab): + print("Building vector cache...") + with open(glove_filename) as f, open(vec_cache_filename, "w") as f2: + for line in f: + tok, vec = line.split(" ", 1) + if tok in vocab: + vocab.remove(tok) + f2.write("{} {}".format(tok, vec)) + +def discrete_tnorm(a, b, tgt_loc, sigma=1, n_steps=100): + def phi(zeta): + return 1 / (np.sqrt(2 * np.pi)) * np.exp(-0.5 * zeta**2) + def Phi(x): + return 0.5 * (1 + erf(x / np.sqrt(2))) + def tgt_loc_update(x): + y1 = phi((a - x) / sigma) + y2 = phi((b - x) / sigma) + x1 = Phi((b - x) / sigma) + x2 = Phi((a - x) / sigma) + denom = x1 - x2 + 1E-4 + return y1 / denom - y2 / denom + + x = tgt_loc + direction = np.sign(tgt_loc - (b - a)) + for _ in range(n_steps): + x = tgt_loc - sigma * tgt_loc_update(x) + tn = truncnorm((a - x) / sigma, (b - x) / sigma, loc=x, scale=sigma) + rrange = np.arange(a, b + 1) + pmf = tn.pdf(rrange) + pmf /= np.sum(pmf) + return pmf + +def discrete_lerp(a, b, ground_truth): + pmf = np.zeros(b - a + 1) + c = int(np.ceil(ground_truth + 1E-8)) + f = int(np.floor(ground_truth)) + pmf[min(c - a, b - a)] = ground_truth - f + pmf[f - a] = c - ground_truth + return pmf + +def smoothed_labels(truth, n_labels): + return discrete_lerp(1, n_labels, truth) + +def preprocess(filename, output_name="sim_sparse.txt"): + print("Preprocessing {}...".format(filename)) + with open(filename) as f: + values = [float(l.strip()) for l in f.readlines()] + values = [" ".join([str(l) for l in smoothed_labels(v, 5)]) for v in values] + with open(os.path.join(os.path.dirname(filename), output_name), "w") as f: + f.write("\n".join(values)) + +def add_vocab(tok_filename, vocab): + with open(tok_filename) as f: + for line in f: + vocab.update(line.strip().split()) + +def main(): + base_conf = data.Configs.base_config() + sick_conf = data.Configs.sick_config() + sick_folder = sick_conf.sick_data + vocab = set() + for name in ("train", "dev", "test"): + preprocess(os.path.join(sick_folder, name, "sim.txt")) + add_vocab(os.path.join(sick_folder, name, "a.toks"), vocab) + add_vocab(os.path.join(sick_folder, name, "b.toks"), vocab) + build_vector_cache(base_conf.wordvecs_file, sick_conf.sick_cache, vocab) + +if __name__ == "__main__": + main() diff --git a/vdpwi/utils/tune.py b/vdpwi/utils/tune.py new file mode 100644 index 0000000..bd4df6c --- /dev/null +++ b/vdpwi/utils/tune.py @@ -0,0 +1,37 @@ +import os +import random + +class RandomParamIterator(object): + def __init__(self, param_sets): + self.param_sets = param_sets + + def random_param_set(self): + param_set = {} + for param_key, param_values in self.param_sets.items(): + param_set[param_key] = random.choice(param_values) + return param_set + +class Tuner(object): + def __init__(self, *iterators, limit=500): + self.iterators = iterators + self.limit = limit + + def start(self): + for i in range(self.limit): + iterator = random.choice(self.iterators) + params = iterator.random_param_set() + print(params) + arg_str = " ".join("--{}={}".format(k, v) for k, v in params.items()) + os.system("python . {} --output_file local_saves/model{}.pt".format(arg_str, i)) + +def main(): + vgg_param_sets = dict(classifer=["vdpwi"], decay=[0.99, 0.95], lr=[5E-4, 1E-4], mbatch_size=[8, 16], + optimizer=["adam", "rmsprop"], weight_decay=[0, 1E-5, 5E-4], momentum=[0, 0.15, 0.05]) + res_param_sets = dict(classifier=["resnet"], clip_norm=[5, 7, 9], decay=[0.9, 0.95], lr=[5E-3, 1E-3, 5E-4], + mbatch_size=[8, 16, 32], rnn_hidden_dim=[150, 250, 300], res_fmaps=[16, 24, 32], res_layers=[4, 8, 16, 24]) + vgg_iterator = RandomParamIterator(vgg_param_sets) + res_iterator = RandomParamIterator(res_param_sets) + Tuner(vgg_iterator).start() + +if __name__ == "__main__": + main() \ No newline at end of file