229 lines
7.5 KiB
Python
229 lines
7.5 KiB
Python
# 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()
|