# -*- coding: utf-8 -*- #for some reason this import has to come first, otherwise gives segfault from transformers import BertTokenizer, BertModel import datetime import argparse import uuid import torch import random import numpy as np import pickle import os import torch.optim as optim import torch.nn as nn from tqdm import tqdm from torch.optim.lr_scheduler import StepLR from utils.util import Utils from torchtext import data from torchtext import datasets from torch.utils.data import Dataset, DataLoader from models.decoder_stgcn import Decoder #from models.navigator_net import Navigator, NavigatorOnlySem, NavigatorOnlySent from data.dataset_sn import SignProdDataset import spacy import nltk from torch import nn import torch.nn.functional as fnn from torch.autograd import Variable from torch.optim import SGD from torchvision.datasets import LSUN from torchvision import transforms from torch.utils.data import Dataset from torchvision.utils import make_grid PAD_IDX = 0 tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") max_input_length = tokenizer.max_model_input_sizes['bert-base-uncased'] nlp = spacy.load("en_core_web_sm") class Navigator(nn.Module): def __init__( self, in_channels, out_channels, depth=1, device = None ): super(Navigator, self).__init__() self.model = nn.Sequential( nn.Linear(1536, 3072), nn.Tanh(), nn.Linear(3072, 3072), nn.Tanh(), nn.Linear(3072, 2304), nn.Tanh(), nn.Linear(2304, 1536) ) def forward(self, sent_feature, sem_feature): x = torch.cat([sent_feature, sem_feature], dim = 1) x = self.model(x) return x def tokenize_and_cut(sentence): tokens = tokenizer.tokenize(sentence) tokens = tokens[:max_input_length-2] return tokens def get_pos(text): doc = nlp(text) pos_list = list() for token in doc: pos_list.append(token.pos_) return pos_list def get_dataset(data_root, batch_size, device, train = True): init_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.cls_token) eos_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.sep_token) pad_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.pad_token) unk_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.unk_token) TEXT = data.Field( batch_first = True, use_vocab = False, tokenize = tokenize_and_cut, preprocessing = tokenizer.convert_tokens_to_ids, init_token = init_token_idx, eos_token = eos_token_idx, pad_token = pad_token_idx, unk_token = unk_token_idx ) RAW = data.RawField() dataset = SignProdDataset( data_root, train, fields = [('src', TEXT), ('kps', RAW), ("sent_feature", RAW), ("sem_feature", RAW), ("pos", RAW), ("aus", RAW), ("z", RAW), ("idx", RAW), ("file_name", RAW)] ) iterator = data.BucketIterator( dataset = dataset, batch_size=batch_size, device=device, sort_key=lambda x: len(x.src), repeat=False, sort=False, shuffle=True if train else False, sort_within_batch= True if train else False ) return iterator def save_ckpt(model, out_fp): torch.save(model.state_dict(), out_fp) def save_Zs(output_root, Z): with open(output_root, "wb") as handler: pickle.dump(Z, handler) def read_pickle(file_fp): with open(file_fp, "rb") as handler: return pickle.load(handler) def read_sentences(dataset_root, files): sentences = list() for idx, file_ in enumerate(files[:50]): file_fp = os.path.join(dataset_root, file_) instance = read_pickle(file_fp) sentences.append((idx, instance["text"])) return sentences def get_k_similar_sentences(pos, sentences, k): sentences_dist = dict() pos_rev_dict = build_pos_rev_dict() for idx, sentence in sentences: pos_s = get_pos(sentence) pos_s_mapped = [pos_rev_dict[p] for p in pos_s] pos_mapped = [pos_rev_dict[p] for p in pos] pos_s_mapped = "".join(pos_s_mapped) pos_mapped = "".join(pos_mapped) distance = nltk.edit_distance(pos_s_mapped, pos_mapped) if distance not in sentences_dist: sentences_dist[distance] = [idx] else: sentences_dist[distance].append(idx) sorted_keys = list(sorted(sentences_dist.keys())) selected_sentences = list() for key in sorted_keys: sentences = sentences_dist[key] for s in sentences: if len(selected_sentences) == k: return selected_sentences selected_sentences.append(s) def build_pos_rev_dict(): return dict( ADJ = "a", ADP = "b", ADV = "c", AUX = "d", CONJ = "e", CCONJ = "f", DET = "g", INTJ = "h", NOUN = "i", NUM = "j", PART = "k", PRON = "l", PROPN = "m", PUNCT = "n", SCONJ = "o", SYM = "p", VERB = "q", X = "r", SPACE = "s" ) def project_l2_ball(z): """ project the vectors in z onto the l2 unit norm ball""" return z / np.maximum(np.sqrt(np.sum(z**2, axis=1))[:, np.newaxis], 1) def validate(test_iterator, navigator, Z, device): epoch_loss = [] navigator.eval() for idx, batch in enumerate(iter(test_iterator)): sent_feature = torch.Tensor(batch.sent_feature).to(device) sem_feature = torch.Tensor(batch.sem_feature).to(device) Zi = torch.Tensor(Z[batch.idx]).to(device).view(-1,768*2) Z_prev = navigator(sent_feature,sem_feature) loss = torch.sum(1 - fnn.cosine_similarity(Zi,Z_prev)) epoch_loss.append(loss.item()) print("Validation loss: {}".format(sum(epoch_loss)/len(epoch_loss))) def main(args): SEED = 1234 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) torch.backends.cudnn.deterministic = True device = torch.device('cuda:{}'.format(args.device)) if args.device != -1 else torch.device('cuda') N_EPOCHS = args.epochs experiment_name = "lr:{}-bs:{}-e:{}-ts:{}".format( args.learning_rate, args.batch_size, args.epochs, datetime.datetime.now() ) #experiment_folder = os.path.join(args.tensorboard_root, experiment_name) #os.makedirs(experiment_folder, exist_ok = True) #tb_writer = SummaryWriter(log_dir = experiment_folder) #encoder = SentiTransformerEncoder().to(device) #decoder = Decoder(device).to(device) #decoder.load_state_dict(torch.load(args.decoder_ckpt)) #decoder.eval() navigator = Navigator(768*2,768*2).to(device) #navigator = NavigatorOnlySem(768*2,768*2).to(device) train_iterator = get_dataset( args.dataset_root, args.batch_size, device, train = True ) test_iterator = get_dataset( args.dataset_root, args.batch_size, device, train = False ) optimizer_n = optim.Adam([ {'params': navigator.parameters(), 'lr': args.learning_rate, 'betas' : (0.5, 0.999)} ]) scheduler_n = StepLR(optimizer_n, step_size = args.save_interval, gamma = 0.5) Z = read_pickle(args.zs_ckpt) mse = nn.MSELoss() step = 0 for epoch in range(N_EPOCHS): epoch_loss = [] for idx, batch in enumerate(iter(train_iterator)): sent_feature = torch.Tensor(batch.sent_feature).to(device) sem_feature = torch.Tensor(batch.sem_feature).to(device) Zi = torch.Tensor(Z[batch.idx]).to(device).view(-1,768*2) optimizer_n.zero_grad() Z_prev = navigator(sent_feature,sem_feature) cosine_distance = torch.sum(1 - fnn.cosine_similarity(Zi,Z_prev)) #mse_distance = mse(Zi, Z_prev) loss = cosine_distance loss.backward() optimizer_n.step() scheduler_n.step() #tb_writer.add_scalar("train/Cosine Distance Loss", loss.item(), step) #tb_writer.add_scalar("train/L2 Loss", mse_distance.item(), step) epoch_loss.append(loss.item()) step = step + 1 if step % args.save_interval == 0 and step > 0: #validate(test_iterator, navigator, Z, device) navigator.train() try: model_folder = os.path.join(args.ckpt_root, experiment_name, "step_{}".format(step)) os.makedirs(model_folder, exist_ok = True) network_fp = os.path.join(model_folder, "navigator.pth") save_ckpt(navigator, network_fp) except: import traceback print(traceback.format_exc()) exit() #encoder.train() print("Epoch loss: {} -- Step: {}".format(sum(epoch_loss)/len(epoch_loss), step)) if __name__ == "__main__": parser = argparse.ArgumentParser(description='Process some integers.') parser.add_argument('--device', '-dev', type = int, default = 0) parser.add_argument('--batch_size', '-bs', type = int, default = 42) parser.add_argument('--epochs', '-e', type = int, default = 10000) parser.add_argument('--save_interval', '-si', type = int, default = 15000) parser.add_argument('--dataset_root', '-d', type = str, default = "/home/rafael/data/how2sign_dataset_speaker1_clean/train/") parser.add_argument('--outputs_root', '-ot', type = str, default = "/home/rafael/data/how2sign_dataset_speaker1_clean/results/") parser.add_argument('--decoder_ckpt', '-dckpt', type = str, default = "/home/rafael/masters/to_move/step_150000_speaker1_h2s/decoder.pth") parser.add_argument('--zs_ckpt', '-zckpt', type = str, default = "/home/rafael/masters/to_move/step_150000_speaker1_h2s/Zs.pkl") parser.add_argument('--learning_rate', '-lr', type = float, default = 0.0001) parser.add_argument('--learning_rate_idenc', '-lridenc', type = float, default = 0.015) parser.add_argument('--learning_rate_dec', '-lrdec', type = float, default = 0.015) parser.add_argument('--learning_rate_disc', '-lrdisc', type = float, default = 0.015) parser.add_argument('--ckpt_root', '-ct', type = str, default = "/home/rafael/masters/to_move/navigator_signer01_h2s_clean") parser.add_argument('--tensorboard_root', '-tr', type = str, default = "/home/rafael/masters/to_move/navigator_signer01_h2s_clean") args = parser.parse_args() main(args)