Spaces:
Running
Running
| # coding: utf-8 | |
| import os | |
| import shutil | |
| import queue | |
| import torch | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| import time | |
| import pickle | |
| import pandas as pd | |
| import numpy as np | |
| import copy | |
| from torch import Tensor | |
| from batch import Batch | |
| from torch.utils.data import Dataset, DataLoader | |
| from torch.nn.utils.rnn import pad_sequence | |
| from helpers import load_config, set_seed, load_checkpoint, log_cfg, make_model_dir, make_logger, ConfigurationError, get_latest_checkpoint, symlink_update | |
| from data import load_data, make_data_iter | |
| from model import build_model, Model | |
| from torch.utils.tensorboard import SummaryWriter | |
| from constants import TARGET_PAD | |
| from loss import Loss | |
| from builders import build_gradient_clipper, build_optimizer, build_scheduler | |
| from plot_videos import plot_video, alter_DTW_timing | |
| from prediction import validate_on_data | |
| class TrainManager: | |
| def __init__(self, model: Model, config: dict, test=False): | |
| train_config = config["training"] | |
| model_dir = train_config["model_dir"] | |
| model_continue = train_config.get("continue", True) | |
| if not os.path.isdir(model_dir): | |
| model_continue = False | |
| if test: | |
| model_continue = True | |
| if test: | |
| self.model_dir = train_config["model_dir"] | |
| else: | |
| self.model_dir = make_model_dir( | |
| model_dir=train_config["model_dir"], | |
| overwrite=train_config.get("overwrite", False) | |
| ) | |
| self.logger = make_logger(model_dir=self.model_dir) | |
| self.logging_freq = train_config.get("logging_freq", 100) | |
| self.valid_report_file = "{}/validations.txt".format(self.model_dir) | |
| self.tb_writer = SummaryWriter(log_dir=self.model_dir + "/tensorboard/") | |
| self.model = model | |
| self.pad_index = self.model.pad_index | |
| self.bos_index = self.model.bos_index | |
| self.target_pad = TARGET_PAD | |
| self.loss = Loss(cfg = config, target_pad=self.target_pad) | |
| self.normalization = "batch" | |
| self.learning_rate_min = train_config.get("learning_rate_min", 1.0e-8) | |
| self.clip_grad_fun = build_gradient_clipper(config=train_config) | |
| self.optimizer = build_optimizer(config=train_config, parameters=model.parameters()) | |
| self.validation_freq = train_config.get("validation_freq", 1000) | |
| self.ckpt_best_queue = queue.Queue(maxsize=train_config.get("keep_last_ckpts", 1)) | |
| self.ckpt_queue = queue.Queue(maxsize=1) | |
| self.eval_metric = train_config.get("eval_metric", "dtw").lower() | |
| if self.eval_metric not in ['bleu', 'chrf', "dtw"]: | |
| raise ConfigurationError("Invalid setting for 'eval_metric', " | |
| "valid options: 'bleu', 'chrf', 'DTW'") | |
| self.early_stopping_metric = train_config.get("early_stopping_metric", | |
| "eval_metric") | |
| if self.early_stopping_metric in ["loss","dtw"]: | |
| self.minimize_metric = True | |
| else: | |
| raise ConfigurationError("Invalid setting for 'early_stopping_metric', " | |
| "valid options: 'loss', 'dtw',.") | |
| self.scheduler, self.scheduler_step_at = build_scheduler( | |
| config=train_config, | |
| scheduler_mode="min" if self.minimize_metric else "max", | |
| optimizer=self.optimizer, | |
| hidden_size=config["model"]["encoder"]["hidden_size"]) | |
| self.level = "word" | |
| self.shuffle = train_config.get("shuffle", True) | |
| self.epochs = train_config["epochs"] | |
| self.batch_size = train_config["batch_size"] | |
| self.batch_type = "sentence" | |
| self.eval_batch_size = train_config.get("eval_batch_size",self.batch_size) | |
| self.eval_batch_type = train_config.get("eval_batch_type",self.batch_type) | |
| self.batch_multiplier = train_config.get("batch_multiplier", 1) | |
| self.max_output_length = train_config.get("max_output_length", None) | |
| self.use_cuda = train_config["use_cuda"] | |
| if self.use_cuda: | |
| self.model.to(device) | |
| self.loss.to(device) | |
| self.steps = 0 | |
| self.stop = False | |
| self.total_tokens = 0 | |
| self.best_ckpt_iteration = 0 | |
| self.best_ckpt_score = np.inf if self.minimize_metric else -np.inf | |
| self.is_best = lambda score: score < self.best_ckpt_score \ | |
| if self.minimize_metric else score > self.best_ckpt_score | |
| self.skip_frames = config["data"].get("skip_frames", 1) | |
| self._log_parameters_list() | |
| if model_continue: | |
| ckpt = get_latest_checkpoint(model_dir) | |
| if ckpt is None: | |
| self.logger.info("Can't find checkpoint in directory %s", ckpt) | |
| else: | |
| self.logger.info("Continuing model from %s", ckpt) | |
| self.init_from_checkpoint(ckpt) | |
| def _log_parameters_list(self) -> None: | |
| model_parameters = filter(lambda p: p.requires_grad, | |
| self.model.parameters()) | |
| n_params = sum([np.prod(p.size()) for p in model_parameters]) | |
| self.logger.info("Total params: %d", n_params) | |
| trainable_params = [n for (n, p) in self.model.named_parameters() | |
| if p.requires_grad] | |
| self.logger.info("Trainable parameters: %s", sorted(trainable_params)) | |
| assert trainable_params | |
| def save_skeleton_files_for_mpje(self, hypotheses, references, inputs, file_paths, data_name, model_dir): | |
| if isinstance(hypotheses, list): | |
| hyp_tensors = [] | |
| for hyp in hypotheses: | |
| if isinstance(hyp, torch.Tensor): | |
| hyp_tensors.append(hyp.cpu()) | |
| else: | |
| hyp_tensors.append(torch.tensor(hyp).cpu()) | |
| hypotheses_tensor = pad_sequence(hyp_tensors, batch_first=True) | |
| else: | |
| hypotheses_tensor = hypotheses.cpu() if hypotheses.is_cuda else hypotheses | |
| if isinstance(references, list): | |
| ref_tensors = [] | |
| for ref in references: | |
| if isinstance(ref, torch.Tensor): | |
| ref_tensors.append(ref.cpu()) | |
| else: | |
| ref_tensors.append(torch.tensor(ref).cpu()) | |
| references_tensor = pad_sequence(ref_tensors, batch_first=True) | |
| else: | |
| references_tensor = references.cpu() if references.is_cuda else references | |
| hyp_path = os.path.join(model_dir, f"{data_name}_hyp_skels.pt") | |
| torch.save(hypotheses_tensor, hyp_path) | |
| ref_path = os.path.join(model_dir, f"{data_name}_ref_skels.pt") | |
| torch.save(references_tensor, ref_path) | |
| return hyp_path, ref_path | |
| def _save_checkpoint(self, type="every") -> None: | |
| model_path = "{}/{}_{}.ckpt".format(self.model_dir, self.steps, type) | |
| state = { | |
| "steps": self.steps, | |
| "total_tokens": self.total_tokens, | |
| "best_ckpt_score": self.best_ckpt_score, | |
| "best_ckpt_iteration": self.best_ckpt_iteration, | |
| "model_state": self.model.state_dict(), | |
| "optimizer_state": self.optimizer.state_dict(), | |
| "scheduler_state": self.scheduler.state_dict() if \ | |
| self.scheduler is not None else None, | |
| } | |
| torch.save(state, model_path) | |
| if type == "best": | |
| if self.ckpt_best_queue.full(): | |
| to_delete = self.ckpt_best_queue.get() | |
| try: | |
| os.remove(to_delete) | |
| except FileNotFoundError: | |
| self.logger.warning("Wanted to delete old checkpoint %s but " | |
| "file does not exist.", to_delete) | |
| self.ckpt_best_queue.put(model_path) | |
| best_path = "{}/best.ckpt".format(self.model_dir) | |
| try: | |
| symlink_update("{}_best.ckpt".format(self.steps), best_path) | |
| except OSError: | |
| torch.save(state, best_path) | |
| elif type == "every": | |
| if self.ckpt_queue.full(): | |
| to_delete = self.ckpt_queue.get() | |
| try: | |
| os.remove(to_delete) | |
| except FileNotFoundError: | |
| self.logger.warning("Wanted to delete old checkpoint %s but " | |
| "file does not exist.", to_delete) | |
| self.ckpt_queue.put(model_path) | |
| every_path = "{}/every.ckpt".format(self.model_dir) | |
| try: | |
| symlink_update("{}_best.ckpt".format(self.steps), every_path) | |
| except OSError: | |
| torch.save(state, every_path) | |
| def init_from_checkpoint(self, path: str) -> None: | |
| model_checkpoint = load_checkpoint(path=path, use_cuda=self.use_cuda) | |
| self.model.load_state_dict(model_checkpoint["model_state"]) | |
| self.optimizer.load_state_dict(model_checkpoint["optimizer_state"]) | |
| if model_checkpoint["scheduler_state"] is not None and \ | |
| self.scheduler is not None: | |
| self.scheduler.load_state_dict(model_checkpoint["scheduler_state"]) | |
| self.steps = model_checkpoint["steps"] | |
| self.total_tokens = model_checkpoint["total_tokens"] | |
| self.best_ckpt_score = model_checkpoint["best_ckpt_score"] | |
| self.best_ckpt_iteration = model_checkpoint["best_ckpt_iteration"] | |
| if torch.cuda.is_available(): | |
| self.model = self.model.cuda() | |
| else: | |
| self.model = self.model.to(torch.device("cpu")) | |
| def train_and_validate(self, train_data: Dataset, valid_data: Dataset) -> None: | |
| train_iter = make_data_iter( | |
| dataset=train_data, | |
| batch_size=self.batch_size, | |
| vocab=self.model.src_vocab, | |
| trg_size=len(self.model.trg_vocab), | |
| shuffle=True, | |
| train=True | |
| ) | |
| val_step = 0 | |
| for epoch_no in range(self.epochs): | |
| self.logger.info("EPOCH %d", epoch_no + 1) | |
| if self.scheduler is not None and self.scheduler_step_at == "epoch": | |
| self.scheduler.step(epoch=epoch_no) | |
| self.model.train() | |
| start = time.time() | |
| total_valid_duration = 0 | |
| start_tokens = self.total_tokens | |
| count = self.batch_multiplier - 1 | |
| epoch_loss = 0 | |
| for batch in iter(train_iter): | |
| self.model.train() | |
| batch = Batch(torch_batch=batch, | |
| pad_index=self.pad_index, | |
| model=self.model) | |
| update = count == 0 | |
| batch_loss = self._train_batch(batch, update=update) | |
| self.tb_writer.add_scalar("train/train_batch_loss", batch_loss, self.steps) | |
| count = self.batch_multiplier if update else count | |
| count -= 1 | |
| epoch_loss += batch_loss.detach().cpu().numpy() | |
| if self.scheduler is not None and self.scheduler_step_at == "step" and update: | |
| self.scheduler.step() | |
| if self.steps % self.logging_freq == 0 and update: | |
| elapsed = time.time() - start - total_valid_duration | |
| elapsed_tokens = self.total_tokens - start_tokens | |
| self.logger.info( | |
| "Epoch %3d Step: %8d Batch Loss: %12.6f " | |
| "Tokens per Sec: %8.0f, Lr: %.6f", | |
| epoch_no + 1, self.steps, batch_loss, | |
| elapsed_tokens / elapsed, | |
| self.optimizer.param_groups[0]["lr"]) | |
| start = time.time() | |
| total_valid_duration = 0 | |
| start_tokens = self.total_tokens | |
| if self.steps % self.validation_freq == 0 and update: | |
| valid_start_time = time.time() | |
| valid_score, valid_loss, valid_references, valid_hypotheses, \ | |
| valid_inputs, all_dtw_scores, valid_file_paths = \ | |
| validate_on_data( | |
| model=self.model, | |
| data=valid_data, | |
| batch_size=self.eval_batch_size, | |
| max_output_length=self.max_output_length, | |
| eval_metric=self.eval_metric, | |
| loss_function=self.loss, | |
| vocab=self.model.src_vocab | |
| ) | |
| self.save_skeleton_files_for_mpje( | |
| hypotheses=valid_hypotheses, | |
| references=valid_references, | |
| inputs=valid_inputs, | |
| file_paths=valid_file_paths, | |
| data_name="dev", | |
| model_dir=self.model_dir | |
| ) | |
| val_step += 1 | |
| self.tb_writer.add_scalar("valid/valid_loss", valid_loss, self.steps) | |
| self.tb_writer.add_scalar("valid/valid_score", valid_score, self.steps) | |
| if self.early_stopping_metric == "loss": | |
| ckpt_score = valid_loss | |
| elif self.early_stopping_metric == "dtw": | |
| ckpt_score = valid_score | |
| else: | |
| ckpt_score = valid_score | |
| new_best = False | |
| self.best = False | |
| if self.is_best(ckpt_score): | |
| self.best = True | |
| self.best_ckpt_score = ckpt_score | |
| self.best_ckpt_iteration = self.steps | |
| self.logger.info( | |
| 'Hooray! New best validation result [%s]!', | |
| self.early_stopping_metric) | |
| if self.ckpt_queue.maxsize > 0: | |
| self.logger.info("Saving new checkpoint.") | |
| new_best = True | |
| self._save_checkpoint(type="best") | |
| display = list(range(0, len(valid_hypotheses), int(np.ceil(len(valid_hypotheses) / 13.15)))) | |
| self.produce_validation_video( | |
| output_joints=valid_hypotheses, | |
| inputs=valid_inputs, | |
| references=valid_references, | |
| model_dir=self.model_dir, | |
| steps=self.steps, | |
| display=display, | |
| type="val_inf", | |
| file_paths=valid_file_paths, | |
| ) | |
| self._save_checkpoint(type="every") | |
| if self.scheduler is not None and self.scheduler_step_at == "validation": | |
| self.scheduler.step(ckpt_score) | |
| self._add_report( | |
| valid_score=valid_score, valid_loss=valid_loss, | |
| eval_metric=self.eval_metric, | |
| new_best=new_best, report_type="val",) | |
| valid_duration = time.time() - valid_start_time | |
| total_valid_duration += valid_duration | |
| self.logger.info( | |
| 'Validation result at epoch %3d, step %8d: Val DTW Score: %6.2f, ' | |
| 'loss: %8.4f, duration: %.4fs', | |
| epoch_no+1, self.steps, valid_score, | |
| valid_loss, valid_duration) | |
| if self.stop: | |
| break | |
| if self.stop: | |
| self.logger.info( | |
| 'Training ended since minimum lr %f was reached.', | |
| self.learning_rate_min) | |
| break | |
| self.logger.info('Epoch %3d: total training loss %.5f', epoch_no+1, | |
| epoch_loss) | |
| else: | |
| self.logger.info('Training ended after %3d epochs.', epoch_no+1) | |
| self.logger.info('Best validation result at step %8d: %6.2f %s.', | |
| self.best_ckpt_iteration, self.best_ckpt_score, | |
| self.early_stopping_metric) | |
| self.tb_writer.close() | |
| def produce_validation_video(self, output_joints, inputs, references, display, model_dir, type, steps="", file_paths=None, dtw_file=None): | |
| if type != "test": | |
| dir_name = model_dir + "/videos/Step_{}/".format(steps) | |
| if not os.path.exists(model_dir + "/videos/"): | |
| os.mkdir(model_dir + "/videos/") | |
| elif type == "test": | |
| dir_name = model_dir + "/test_videos/" | |
| if not os.path.exists(dir_name): | |
| os.mkdir(dir_name) | |
| for i in display: | |
| seq = output_joints[i] | |
| ref_seq = references[i] | |
| input = inputs[i] | |
| gloss_label = input[0] | |
| if input[1] != "</s>": | |
| gloss_label += "_" + input[1] | |
| if input[2] != "</s>": | |
| gloss_label += "_" + input[2] | |
| timing_hyp_seq, ref_seq_count, dtw_score = alter_DTW_timing(seq, ref_seq) | |
| video_ext = "{}_{}.mp4".format(gloss_label, "{0:.2f}".format(float(dtw_score)).replace(".", "_")) | |
| if file_paths is not None: | |
| sequence_ID = file_paths[i] | |
| else: | |
| sequence_ID = None | |
| print(sequence_ID + ' dtw: ' + '{0:.2f}'.format(float(dtw_score))) | |
| if dtw_file != None: | |
| dtw_file.writelines(sequence_ID + ' ' + '{0:.2f}'.format(float(dtw_score)) + '\n') | |
| plot_video(joints=timing_hyp_seq, | |
| file_path=dir_name, | |
| video_name=video_ext, | |
| references=ref_seq_count, | |
| skip_frames=self.skip_frames, | |
| sequence_ID=sequence_ID) | |
| def save_skels(self, output_joints, display, model_dir, type, file_paths=None): | |
| picklefile = open(model_dir + "/phoenix14t.skels.%s" % type, "wb") | |
| try: | |
| csvIn = pd.read_csv(model_dir + "/csv/%s_phoenix2014t.csv" % type, sep='|',encoding='utf-8') | |
| except FileNotFoundError: | |
| print(f"Warning: CSV file not found at {model_dir}/csv/{type}_phoenix2014t.csv") | |
| print("Saving skeleton data without CSV metadata...") | |
| pickle_list = [] | |
| for i in display: | |
| name = file_paths[i] if file_paths else f"sequence_{i}" | |
| seq = output_joints[i].cpu()[:,:-1] | |
| sign = seq.clone().detach().to(torch.float32) | |
| dict_num = {'name': name, 'signer': 'unknown', 'gloss': 'unknown', 'text': 'unknown', 'sign': sign} | |
| pickle_list.append(dict_num) | |
| pickle.dump(pickle_list, picklefile) | |
| print("The skeletons of %s data have been saved without metadata." % type) | |
| return | |
| pickle_list = [] | |
| for i in display: | |
| name = file_paths[i] | |
| video = os.path.basename(name) | |
| matching_rows = csvIn[csvIn['id'] == video] | |
| if matching_rows.empty: | |
| print(f"Warning: No exact match for video id={video}") | |
| partial_matches = csvIn[csvIn['id'].str.contains(video.split('_')[0], na=False)] | |
| if not partial_matches.empty: | |
| signer = partial_matches.iloc[0]['signer'] | |
| gloss = partial_matches.iloc[0]['annotation'] | |
| text = partial_matches.iloc[0]['translation'] | |
| else: | |
| signer = 'unknown' | |
| gloss = 'unknown' | |
| text = 'unknown' | |
| else: | |
| signer = matching_rows.iloc[0]['signer'] | |
| gloss = matching_rows.iloc[0]['annotation'] | |
| text = matching_rows.iloc[0]['translation'] | |
| seq = output_joints[i].cpu()[:,:-1] | |
| sign = seq.clone().detach().to(torch.float32) | |
| dict_num = {'name': name, 'signer': signer, 'gloss': gloss, 'text': text, 'sign': sign} | |
| pickle_list.append(dict_num) | |
| pickle.dump(pickle_list, picklefile) | |
| print("The skeletons of %s data have been saved." % type) | |
| def _train_batch(self, batch: Batch, update: bool = True) -> Tensor: | |
| batch_loss = self.model.get_loss_for_batch( | |
| is_train=True, | |
| batch=batch, | |
| loss_function=self.loss | |
| ) | |
| if self.normalization == "batch": | |
| normalizer = batch.nseqs | |
| elif self.normalization == "tokens": | |
| normalizer = batch.ntokens | |
| else: | |
| raise NotImplementedError("Only normalize by 'batch' or 'tokens'") | |
| norm_batch_loss = batch_loss / normalizer | |
| norm_batch_multiply = norm_batch_loss / self.batch_multiplier | |
| norm_batch_multiply.backward() | |
| if self.clip_grad_fun is not None: | |
| self.clip_grad_fun(params=self.model.parameters()) | |
| if update: | |
| self.optimizer.step() | |
| self.optimizer.zero_grad() | |
| self.steps += 1 | |
| self.total_tokens += batch.ntokens | |
| return norm_batch_loss | |
| def _add_report(self, valid_score: float, valid_loss: float, eval_metric: str, new_best: bool = False, report_type: str = "val") -> None: | |
| current_lr = -1 | |
| for param_group in self.optimizer.param_groups: | |
| current_lr = param_group['lr'] | |
| if current_lr < self.learning_rate_min: | |
| self.stop = True | |
| if report_type == "val": | |
| with open(self.valid_report_file, 'a') as opened_file: | |
| opened_file.write( | |
| "Steps: {} Loss: {:.5f}| DTW: {:.3f}|" | |
| " LR: {:.6f} {}\n".format( | |
| self.steps, valid_loss, valid_score, | |
| current_lr, "*" if new_best else "")) | |
| def train(cfg_file: str, ckpt=None): | |
| cfg = load_config(cfg_file) | |
| set_seed(seed=cfg["training"].get("random_seed", 42)) | |
| train_data, dev_data, test_data, src_vocab, trg_vocab = load_data(cfg=cfg) | |
| model = build_model(cfg=cfg, src_vocab=src_vocab, trg_vocab=trg_vocab) | |
| if model is None: | |
| raise ValueError("Model creation failed. Check your build_model function.") | |
| if ckpt is not None: | |
| use_cuda = cfg["training"].get("use_cuda", True) | |
| if not os.path.isfile(ckpt): | |
| raise FileNotFoundError(f"Checkpoint file not found: {ckpt}") | |
| checkpoint = load_checkpoint(ckpt, use_cuda=use_cuda) | |
| if checkpoint is None or "model_state" not in checkpoint: | |
| raise ValueError(f"Failed to load model state from checkpoint: {ckpt}") | |
| state_dict = checkpoint["model_state"] | |
| if "src_embed.lut.weight" in state_dict: | |
| state_dict.pop("src_embed.lut.weight") | |
| model.load_state_dict(state_dict, strict=False) | |
| model.to(device) | |
| trainer = TrainManager(config=cfg, model=model, test=False) | |
| shutil.copy2(cfg_file, os.path.join(trainer.model_dir, "Sign-IDD.yaml")) | |
| log_cfg(cfg, trainer.logger) | |
| trainer.train_and_validate(train_data=train_data, valid_data=dev_data) | |
| def test(cfg_file: str, ckpt: str = None): | |
| cfg = load_config(cfg_file) | |
| model_dir = cfg["training"]["model_dir"] | |
| if ckpt is None: | |
| ckpt = get_latest_checkpoint(model_dir, post_fix="_best") | |
| if ckpt is None: | |
| raise FileNotFoundError(f"No checkpoint found in directory: {model_dir}") | |
| train_data, dev_data, test_data, src_vocab, trg_vocab = load_data(cfg=cfg) | |
| model = build_model(cfg=cfg, src_vocab=src_vocab, trg_vocab=trg_vocab) | |
| if model is None: | |
| raise ValueError("Model creation failed. Check your build_model function.") | |
| if not os.path.isfile(ckpt): | |
| raise FileNotFoundError(f"Checkpoint file not found: {ckpt}") | |
| checkpoint = load_checkpoint(ckpt, use_cuda=cfg["training"].get("use_cuda", True)) | |
| if checkpoint is None or "model_state" not in checkpoint: | |
| raise ValueError(f"Failed to load model state from checkpoint: {ckpt}") | |
| state_dict = checkpoint["model_state"] | |
| emb_key = "src_embed.lut.weight" | |
| model_emb = None | |
| if hasattr(model, "src_embed"): | |
| src_embed_obj = getattr(model, "src_embed") | |
| if hasattr(src_embed_obj, "lut") and hasattr(src_embed_obj.lut, "weight"): | |
| model_emb = src_embed_obj.lut.weight | |
| elif hasattr(src_embed_obj, "weight"): | |
| model_emb = src_embed_obj.weight | |
| if emb_key in state_dict: | |
| ckpt_shape = state_dict[emb_key].shape | |
| model_shape = model_emb.shape if model_emb is not None else None | |
| if model_shape is None or ckpt_shape != model_shape: | |
| print(f"[INFO] Skipping {emb_key} due to shape mismatch: checkpoint {ckpt_shape} vs model {model_shape}") | |
| state_dict.pop(emb_key) | |
| else: | |
| pass | |
| model.load_state_dict(state_dict, strict=False) | |
| use_cuda_flag = cfg["training"].get("use_cuda", True) | |
| device_local = torch.device("cuda" if use_cuda_flag and torch.cuda.is_available() else "cpu") | |
| model.to(device_local) | |
| cfg_no_continue = copy.deepcopy(cfg) | |
| cfg_no_continue["training"] = dict(cfg_no_continue.get("training", {})) | |
| cfg_no_continue["training"]["continue"] = False | |
| trainer = TrainManager(model=model, config=cfg_no_continue, test=True) | |
| data_to_predict = {"dev": dev_data, "test": test_data} | |
| batch_size = cfg["training"].get("eval_batch_size", cfg["training"]["batch_size"]) | |
| batch_type = cfg["training"].get("eval_batch_type", cfg["training"].get("batch_type", "sentence")) | |
| eval_metric = cfg["training"]["eval_metric"] | |
| max_output_length = cfg["training"].get("max_output_length", None) | |
| for data_name, dataset in data_to_predict.items(): | |
| score, loss, references, hypotheses, inputs, all_dtw_scores, file_paths = validate_on_data( | |
| model=model, | |
| data=dataset, | |
| batch_size=batch_size, | |
| max_output_length=max_output_length, | |
| eval_metric=eval_metric, | |
| loss_function=None, | |
| vocab=model.src_vocab | |
| ) | |
| trainer.save_skeleton_files_for_mpje( | |
| hypotheses=hypotheses, | |
| references=references, | |
| inputs=inputs, | |
| file_paths=file_paths, | |
| data_name=data_name, | |
| model_dir=model_dir | |
| ) | |
| output_dir = os.path.join(model_dir, "test_videos") | |
| os.makedirs(output_dir, exist_ok=True) | |
| dtw_path = os.path.join(output_dir, f"{data_name}_dtw.txt") | |
| with open(dtw_path, "w") as f: | |
| f.write(f"DTW Score of {data_name} set: {score:.3f}\n") | |
| print(f"DTW Score of {data_name} set: {score:.3f}") | |
| display = list(range(len(hypotheses))) | |
| trainer.save_skels(output_joints=hypotheses, display=display, model_dir=model_dir, type=data_name, file_paths=file_paths) | |
| with open(dtw_path, "a") as dtw_file: | |
| trainer.produce_validation_video( | |
| output_joints=hypotheses, | |
| inputs=inputs, | |
| references=references, | |
| model_dir=model_dir, | |
| display=display, | |
| type="test", | |
| file_paths=file_paths, | |
| dtw_file=dtw_file, | |
| ) | |
| return model_dir |