meng2003/music2dance

1import torch2from contextlib import contextmanager3from collections import OrderedDict4print("HOOOOOO")5from pytorch_lightning import LightningModule6print("HOOOOOO")7from .optimizer import get_scheduler, get_optimizers8 9from models.util.generation import autoregressive_generation_multimodal10 11# Benefits of having one skeleton, e.g. for train - is that you can keep all the incremental changes in12# one single code, making it your streamlined and updated script -- no need to keep separate logs on how13# to implement stuff14 15class BaseModel(LightningModule):16    def __init__(self, opt):17        super().__init__()18        self.save_hyperparameters(vars(opt))19        self.opt = opt20        self.parse_base_arguments()21        self.optimizers = []22        self.schedulers = []23 24    def parse_base_arguments(self):25        # import pdb;pdb.set_trace()26        self.input_mods = str(self.opt.input_modalities).split(",")27        self.output_mods = str(self.opt.output_modalities).split(",")28        self.dins = [int(x) for x in str(self.opt.dins).split(",")]29        self.douts = [int(x) for x in str(self.opt.douts).split(",")]30        self.input_lengths = [int(x) for x in str(self.opt.input_lengths).split(",")]31        self.output_lengths = [int(x) for x in str(self.opt.output_lengths).split(",")]32        self.output_time_offsets = [int(x) for x in str(self.opt.output_time_offsets).split(",")]33        self.input_time_offsets = [int(x) for x in str(self.opt.input_time_offsets).split(",")]34 35        if len(self.output_time_offsets) < len(self.output_mods):36            if len(self.output_time_offsets) == 1:37                self.output_time_offsets = self.output_time_offsets*len(self.output_mods)38            else:39                raise Exception("number of output_time_offsets doesnt match number of output_mods")40 41        if len(self.input_time_offsets) < len(self.input_mods):42            if len(self.input_time_offsets) == 1:43                self.input_time_offsets = self.input_time_offsets*len(self.input_mods)44            else:45                raise Exception("number of input_time_offsets doesnt match number of input_mods")46 47        input_mods = self.input_mods48        if self.opt.input_types is None:49            input_types = ["c" for inp in input_mods]50        else:51            input_types = self.opt.input_types.split(",")52 53        if self.opt.input_fix_length_types is None:54            input_fix_length_types = ["end" for inp in input_mods]55        else:56            input_fix_length_types = self.opt.input_fix_length_types.split(",")57 58        if self.opt.output_fix_length_types is None:59            output_fix_length_types = ["end" for inp in input_mods]60        else:61            output_fix_length_types = self.opt.output_fix_length_types.split(",")62 63        #fix_length_types_dict = {mod:output_fix_length_types[i] for i,mod in enumerate(output_mods)}64        #fix_length_types_dict.update({mod:input_fix_length_types[i] for i,mod in enumerate(input_mods)})65 66        assert len(input_types) == len(input_mods)67        assert len(input_fix_length_types) == len(input_mods)68        assert len(output_fix_length_types) == len(input_mods)69        self.input_types = input_types70        self.input_fix_length_types = input_fix_length_types71        self.output_fix_length_types = output_fix_length_types72 73        if self.opt.input_num_tokens is None:74            self.input_num_tokens = [0 for inp in self.input_mods]75        else:76            self.input_num_tokens  = [int(x) for x in self.opt.input_num_tokens.split(",")]77 78        if self.opt.output_num_tokens is None:79            self.output_num_tokens = [0 for inp in self.output_mods]80        else:81            self.output_num_tokens  = [int(x) for x in self.opt.output_num_tokens.split(",")]82 83 84    def name(self):85        return 'BaseModel'86 87    #def setup_opt(self, is_train):88    #    pass89 90    def configure_optimizers(self):91        optimizers = get_optimizers(self, self.opt)92        schedulers = [get_scheduler(optimizer, self.opt) for optimizer in self.optimizers]93        return optimizers, schedulers94        #return self.optimizers95 96    def set_inputs(self, data):97        # BTC -> TBC98        self.inputs = []99        self.targets = []100        for i, mod in enumerate(self.input_mods):101            input_ = data["in_"+mod]102            input_ = input_.permute(1,0,2)103            self.inputs.append(input_)104        for i, mod in enumerate(self.output_mods):105            target_ = data["out_"+mod]106            target_ = target_.permute(1,0,2)107            self.targets.append(target_)108 109    def generate(self,features, teacher_forcing=False, ground_truth=False):110        output_seq = autoregressive_generation_multimodal(features, self, autoreg_mods=self.output_mods, teacher_forcing=teacher_forcing, ground_truth=ground_truth)111        return output_seq112 113    # modify parser to add command line options,114    # and also change the default values if needed115    @staticmethod116    def modify_commandline_options(parser, is_train):117        """118        ABSTRACT METHOD119        :param parser:120        :param is_train:121        :return:122        """123        return parser124 125    def test_step(self, batch, batch_idx):126        self.eval()127        loss = self.training_step(batch, batch_idx)128        # print(loss)129        return {"test_loss": loss}130 131    def test_epoch_end(self, outputs):132        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()133        logs = {'test_loss': avg_loss}134 135        return {'log': logs}136 137