Upload ./RepCodec/train.py with huggingface_hub

RepCodec/train.py

@@ -0,0 +1,228 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import argparse
+import logging
+
+import os
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+)
+logger = logging.getLogger("repcodec_train")  # init logger before other modules
+
+import random
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+from dataloader import ReprDataset, ReprCollater
+from losses.repr_reconstruct_loss import ReprReconstructLoss
+from repcodec.RepCodec import RepCodec
+from trainer.autoencoder import Trainer
+
+
+class TrainMain:
+    def __init__(self, args):
+        # Fix seed and make backends deterministic
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if not torch.cuda.is_available():
+            self.device = torch.device('cpu')
+            logger.info(f"device: cpu")
+        else:
+            self.device = torch.device('cuda:0')  # only supports single gpu for now
+            logger.info(f"device: gpu")
+            torch.cuda.manual_seed_all(args.seed)
+            if args.disable_cudnn == "False":
+                torch.backends.cudnn.benchmark = True
+
+        # initialize config
+        with open(args.config, 'r') as f:
+            self.config = yaml.load(f, Loader=yaml.FullLoader)
+        self.config.update(vars(args))
+
+        # initialize model folder
+        expdir = os.path.join(args.exp_root, args.tag)
+        os.makedirs(expdir, exist_ok=True)
+        self.config["outdir"] = expdir
+
+        # save config
+        with open(os.path.join(expdir, "config.yml"), "w") as f:
+            yaml.dump(self.config, f, Dumper=yaml.Dumper)
+        for key, value in self.config.items():
+            logger.info(f"{key} = {value}")
+
+        # initialize attribute
+        self.resume: str = args.resume
+        self.data_loader = None
+        self.model = None
+        self.optimizer = None
+        self.scheduler = None
+        self.criterion = None
+        self.trainer = None
+
+        # initialize batch_length
+        self.batch_length: int = self.config['batch_length']
+        self.data_path: str = self.config['data']['path']
+
+    def initialize_data_loader(self):
+        train_set = self._build_dataset("train")
+        valid_set = self._build_dataset("valid")
+        collater = ReprCollater()
+
+        logger.info(f"The number of training files = {len(train_set)}.")
+        logger.info(f"The number of validation files = {len(valid_set)}.")
+        dataset = {"train": train_set, "dev": valid_set}
+        self._set_data_loader(dataset, collater)
+
+    def define_model_optimizer_scheduler(self):
+        # model arch
+        self.model = {
+            "repcodec": RepCodec(**self.config["model_params"]).to(self.device)
+        }
+        logger.info(f"Model Arch:\n{self.model['repcodec']}")
+
+        # opt
+        optimizer_class = getattr(
+            torch.optim,
+            self.config["model_optimizer_type"]
+        )
+        self.optimizer = {
+            "repcodec": optimizer_class(
+                self.model["repcodec"].parameters(),
+                **self.config["model_optimizer_params"]
+            )
+        }
+
+        # scheduler
+        scheduler_class = getattr(
+            torch.optim.lr_scheduler,
+            self.config.get("model_scheduler_type", "StepLR"),
+        )
+        self.scheduler = {
+            "repcodec": scheduler_class(
+                optimizer=self.optimizer["repcodec"],
+                **self.config["model_scheduler_params"]
+            )
+        }
+
+    def define_criterion(self):
+        self.criterion = {
+            "repr_reconstruct_loss": ReprReconstructLoss(
+                **self.config.get("repr_reconstruct_loss_params", {}),
+            ).to(self.device)
+        }
+
+    def define_trainer(self):
+        self.trainer = Trainer(
+            steps=0,
+            epochs=0,
+            data_loader=self.data_loader,
+            model=self.model,
+            criterion=self.criterion,
+            optimizer=self.optimizer,
+            scheduler=self.scheduler,
+            config=self.config,
+            device=self.device
+        )
+
+    def initialize_model(self):
+        initial = self.config.get("initial", "")
+        if os.path.exists(self.resume):  # resume from trained model
+            self.trainer.load_checkpoint(self.resume)
+            logger.info(f"Successfully resumed from {self.resume}.")
+        elif os.path.exists(initial):  # initial new model with the pre-trained model
+            self.trainer.load_checkpoint(initial, load_only_params=True)
+            logger.info(f"Successfully initialize parameters from {initial}.")
+        else:
+            logger.info("Train from scrach")
+
+    def run(self):
+        assert self.trainer is not None
+        self.trainer: Trainer
+        try:
+            logger.info(f"The current training step: {self.trainer.steps}")
+            self.trainer.train_max_steps = self.config["train_max_steps"]
+            if not self.trainer._check_train_finish():
+                self.trainer.run()
+        finally:
+            self.trainer.save_checkpoint(
+                os.path.join(self.config["outdir"], f"checkpoint-{self.trainer.steps}steps.pkl")
+            )
+            logger.info(f"Successfully saved checkpoint @ {self.trainer.steps}steps.")
+
+    def _build_dataset(
+            self, subset: str
+    ) -> ReprDataset:
+        data_dir = os.path.join(
+            self.data_path, self.config['data']['subset'][subset]
+        )
+        params = {
+            "data_dir": data_dir,
+            "batch_len": self.batch_length
+        }
+        return ReprDataset(**params)
+
+    def _set_data_loader(self, dataset, collater):
+        self.data_loader = {
+            "train": DataLoader(
+                dataset=dataset["train"],
+                shuffle=True,
+                collate_fn=collater,
+                batch_size=self.config["batch_size"],
+                num_workers=self.config["num_workers"],
+                pin_memory=self.config["pin_memory"],
+            ),
+            "dev": DataLoader(
+                dataset=dataset["dev"],
+                shuffle=False,
+                collate_fn=collater,
+                batch_size=self.config["batch_size"],
+                num_workers=0,
+                pin_memory=False,  # save some memory. set to True if you have enough memory.
+            ),
+        }
+
+
+def train():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c", "--config", type=str, required=True,
+        help="the path of config yaml file."
+    )
+    parser.add_argument(
+        "--tag", type=str, required=True,
+        help="the outputs will be saved to exp_root/tag/"
+    )
+    parser.add_argument(
+        "--exp_root", type=str, default="exp"
+    )
+    parser.add_argument(
+        "--resume", default="", type=str, nargs="?",
+        help='checkpoint file path to resume training. (default="")',
+    )
+    parser.add_argument("--seed", default=1337, type=int)
+    parser.add_argument("--disable_cudnn", choices=("True", "False"), default="False", help="Disable CUDNN")
+    args = parser.parse_args()
+
+    train_main = TrainMain(args)
+    train_main.initialize_data_loader()
+    train_main.define_model_optimizer_scheduler()
+    train_main.define_criterion()
+    train_main.define_trainer()
+    train_main.initialize_model()
+    train_main.run()
+
+
+if __name__ == '__main__':
+    train()