main.py

# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import fire
import pprint
import yaml
import torch
import wandb
import logging
import popart
import poptorch
import numpy as np
from src.conformer_encoder import ConformerEncoder
from src.transformer_decoder import TransformerDecoder
from src.global_cmvn import GlobalMVN
from src.label_smoothing_loss import LabelSmoothingLoss, CrossEntropyLoss
from src.conformer import Conformer
from src.utils.initializer import initialize
from src.trainer import Trainer
from src.utils.lr_scheduler import WarmupLR
from src.utils.checkpoint import CheckPoint
from src.iterator.dataset import Dataset
from src.utils.file_utils import read_symbol_table, read_non_lang_symbols
from src.iterator.dataset import IPUCollateFn
from src.iterator.generate_data import GenerateDataset
from src.utils.cmvn import load_cmvn

import copy
import popdist
import popdist.poptorch
import horovod.torch as hvd
import os


def load_yaml(yaml_file):
    args = yaml.safe_load(open(yaml_file, "r"))
    return args


class Workflow:
    def __init__(self, config_file="configs/train.yaml", **kwargs):
        self.config_file = config_file
        self.args = load_yaml(self.config_file)
        self.parse_args(kwargs)
        self.ipu_options = None
        self.vocab_size = self.args["decoder"]["vocab_size"]
        self.train_dataset = None
        self.val_dataset = None
        self.train_iterator = None
        self.val_iterator = None
        self.model = None
        self.trainer = None
        self.checkpoint = None
        self.optimizer = None
        self.scheduler = None
        self.wandb = None

    def parse_args(self, kwargs):
        for key, value in kwargs.items():
            try:
                module_name, config_name = key.split(".")
                self.args[module_name][config_name] = value
            except:
                raise KeyError(f"{key} is not found in {self.config_file}.")

    def build_optimizer(self):
        accum_type = self.dtype
        self.optimizer = poptorch.optim.Adam(self.model.parameters(), accum_type=accum_type, **self.args["optimizer"])
        self.scheduler = WarmupLR(self.optimizer, **self.args["scheduler"])

    def build_wandb(self):
        if not self.args["trainer"]["wandb_project_name"]:
            self.wandb = None
        else:
            if self.args["popdist_rank"] == 0:
                wandb.init(
                    project=self.args["trainer"]["wandb_project_name"],
                    name=self.args["trainer"]["wandb_run_name"],
                    settings=wandb.Settings(console="off"),
                )
                wandb.config.update(self.args)
                self.wandb = wandb

    def build_logger(self):
        """build logger both for print in terminal and log.txt"""
        logger_args = self.args["trainer"]["logger"]
        log_level = {"info": logging.INFO, "debug": logging.DEBUG}.get(logger_args["level"], logging.INFO)
        self.logger = logging.getLogger(logger_args["name"])
        self.logger.setLevel(level=log_level)
        formatter = logging.Formatter("%(asctime)s %(filename)s line:%(lineno)d %(levelname)s %(message)s")
        handler = logging.FileHandler(logger_args["log_file"])
        handler.setFormatter(formatter)
        console = logging.StreamHandler()
        console.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.addHandler(console)

    def launch_popdist(self):
        if popdist.isPopdistEnvSet():
            hvd.init()
            self.args["popdist_replicas"] = int(popdist.getNumLocalReplicas())
            self.args["popdist_rank"] = popdist.getInstanceIndex()
            self.args["NumInstances"] = popdist.getNumInstances()
        else:
            self.args["NumInstances"] = 1
            self.args["popdist_rank"] = 0

    def build_ipu_options(self, is_train=True):
        ipu_args = self.args["ipu_options"]
        ipus_per_replica = len(ipu_args["pipeline"]) + 1
        replicas = ipu_args["num_replicas"]

        if is_train:
            lbs = self.args["train_iterator"]["batch_size"]
            ga = ipu_args["gradient_accumulation"]
            if popdist.isPopdistEnvSet():
                self.ipu_options = popdist.poptorch.Options(ipus_per_replica=ipus_per_replica)
            else:
                self.ipu_options = poptorch.Options()
                self.ipu_options.replicationFactor(replicas)
            self.ipu_options.randomSeed(self.args["trainer"]["random_seed"])
            self.ipu_options.autoRoundNumIPUs(True)
            self.ipu_options.Training.gradientAccumulation(ga)
            self.ipu_options.deviceIterations(ipu_args["device_iterations"])
            self.ipu_options.outputMode(poptorch.OutputMode.Final)
            self.ipu_options.enableExecutableCaching(ipu_args["executable_cache_dir"])
            self.ipu_options.Training.accumulationAndReplicationReductionType(poptorch.ReductionType.Mean)
            self.ipu_options.setExecutionStrategy(poptorch.PipelinedExecution(poptorch.AutoStage.AutoIncrement))
            self.ipu_options.TensorLocations.setOptimizerLocation(
                poptorch.TensorLocationSettings()
                .useOnChipStorage(not ipu_args["optimizer_state_offchip"])
                .useReplicatedTensorSharding(ipu_args["replicated_tensor_sharding"] if replicas > 1 else False)
            )
            if ipu_args["enable_half_partials"]:
                self.ipu_options.Precision.setPartialsType(torch.half)
            self.ipu_options.setAvailableMemoryProportion(
                {f"IPU{i}": ipu_args["available_memory_proportion"][i] for i in range(len(ipu_args["pipeline"]) + 1)}
            )

            self.ipu_options.Precision.enableStochasticRounding(ipu_args["enable_stochastic_rounding"])
            # PopART settings
            # enable recomputation in pipeline mode
            self.ipu_options._Popart.set("disableGradAccumulationTensorStreams", True)
            self.ipu_options._Popart.set("outlineThreshold", 10.0)
            self.ipu_options._Popart.set("timeLimitScheduler", float(120))
            self.ipu_options._Popart.set("accumulateOuterFragmentSettings.excludedVirtualGraphs", ["0"])
            self.ipu_options._Popart.set("scheduleNonWeightUpdateGradientConsumersEarly", True)
            self.ipu_options._Popart.setPatterns(
                {"TiedGather": True, "TiedGatherAccumulate": True, "UpdateInplacePrioritiesForIpu": True}
            )

        else:
            lbs = self.args["val_iterator"]["batch_size"]
            ga = 1
            self.ipu_options = poptorch.Options()
            self.ipu_options.replicationFactor(replicas)
            self.ipu_options.autoRoundNumIPUs(True)
            self.ipu_options.deviceIterations(ipu_args["device_iterations"])
            self.ipu_options.Training.gradientAccumulation(ga)
            self.ipu_options.outputMode(poptorch.OutputMode.Final)
            self.ipu_options.setExecutionStrategy(poptorch.PipelinedExecution(poptorch.AutoStage.SameAsIpu))
            if ipu_args["enable_half_partials"]:
                self.ipu_options.Precision.setPartialsType(torch.half)
            self.ipu_options.setAvailableMemoryProportion(
                {f"IPU{i}": ipu_args["available_memory_proportion"][i] for i in range(len(ipu_args["pipeline"]) + 1)}
            )

        if ipu_args["enable_profiling"]:
            ampstr = ",".join([str(amp_) for amp_ in ipu_args["available_memory_proportion"]])
            sdk_version = "-".join(os.environ.get("POPLAR_SDK_ENABLED").split("/")[-2].split("-")[-3:])
            pipeline_str = ""
            for stage in ipu_args["pipeline"]:
                # encoder and decoder
                prefix = stage[0].split("__")
                if len(prefix) == 1:
                    pipeline_str += prefix[0][:3]  # eg:ctc
                else:
                    pipeline_str += prefix[0][:3] + prefix[1]  # eg:enc3

            profile_path = os.path.join(
                os.path.abspath(ipu_args["profile_path"]),
                f"bs{lbs}-ga{ga}-amp{ampstr}-rep{replicas}-pl{ipus_per_replica}-{pipeline_str}-{sdk_version}",
            )

            os.makedirs(profile_path, exist_ok=True)

            engine_options = {
                "autoReport.directory": profile_path,
                "target.syncReplicasIndependently": "true",
                "debug.allowOutOfMemory": "true",
                "profiler.includeFlopEstimates": "true",
                "profiler.includeCycleEstimates": "true",
                "autoReport.all": "true",
                "autoReport.executionProfileProgramRunCount": "2",
            }
            self.ipu_options._Popart.set("engineOptions", engine_options)

        if ipu_args["compile_only"]:
            self.ipu_options.useOfflineIpuTarget()

    def set_random_seed(self):
        np.random.seed(self.args["trainer"]["random_seed"])
        torch.manual_seed(self.args["trainer"]["random_seed"])

    def build_checkpoints(self):
        self.checkpoint = CheckPoint(0, self.logger)

    def build_dataset(self):
        self.train_conf = self.args["train_conf"]
        self.cv_conf = copy.deepcopy(self.train_conf)
        self.cv_conf["speed_perturb"] = False
        self.cv_conf["spec_aug"] = False
        self.cv_conf["spec_sub"] = False
        self.cv_conf["shuffle"] = False

        if not self.args["train_dataset"]["use_generated_data"]:
            lang = self.args["vocab"]["vocab_path"]
            symbol_table = read_symbol_table(lang)
            self.train_dataset = Dataset(
                self.args["train_dataset"]["data_mode"],
                self.args["train_dataset"]["data_list"],
                symbol_table,
                self.train_conf,
                None,
                None,
                True,
            )
            self.val_dataset = Dataset(
                self.args["train_dataset"]["data_mode"],
                self.args["val_dataset"]["data_list"],
                symbol_table,
                self.cv_conf,
                None,
                None,
                partition=False,
            )
        else:
            self.train_dataset = GenerateDataset(self.args)

    def build_iterator(self, is_train=True):
        if not self.ipu_options:
            self.build_ipu_options(is_train)
        self.build_dataset()
        if self.args["train_iterator"]["async_mode"]:
            mode = poptorch.DataLoaderMode.Async
        else:
            mode = poptorch.DataLoaderMode.Sync
        if not self.args["train_dataset"]["use_generated_data"]:
            self.train_iterator = poptorch.DataLoader(
                self.ipu_options,
                self.train_dataset,
                batch_size=self.args["train_iterator"]["batch_size"],
                num_workers=self.args["train_iterator"]["num_workers"],
                persistent_workers=self.args["train_iterator"]["persistent_workers"],
                async_options=self.args["train_iterator"]["async_options"],
                mode=mode,
                collate_fn=IPUCollateFn(
                    self.train_conf["filter_conf"]["max_length"],
                    self.train_conf["filter_conf"]["token_max_length"],
                    dtype=self.dtype,
                    sos_id=self.vocab_size - 1,
                    eos_id=self.vocab_size - 1,
                ),
            )
            self.val_iterator = poptorch.DataLoader(
                self.ipu_options,
                self.val_dataset,
                batch_size=self.args["val_iterator"]["batch_size"],
                num_workers=self.args["val_iterator"]["num_workers"],
                persistent_workers=self.args["train_iterator"]["persistent_workers"],
                async_options=self.args["train_iterator"]["async_options"],
                mode=mode,
                collate_fn=IPUCollateFn(
                    self.train_conf["filter_conf"]["max_length"],
                    self.train_conf["filter_conf"]["token_max_length"],
                    dtype=self.dtype,
                    sos_id=self.vocab_size - 1,
                    eos_id=self.vocab_size - 1,
                ),
            )
        else:
            self.train_iterator = poptorch.DataLoader(
                self.ipu_options,
                dataset=self.train_dataset,
                persistent_workers=self.args["train_iterator"]["persistent_workers"],
                async_options=self.args["train_iterator"]["async_options"],
                mode=mode,
                shuffle=False,
                batch_size=self.args["train_iterator"]["batch_size"],
                num_workers=self.args["train_iterator"]["num_workers"],
            )

    def build_model(self):
        self.dtype = torch.float16 if self.args["train_dataset"]["dtype"] == "FLOAT16" else torch.float32
        encoder = ConformerEncoder(dtype=self.dtype, **self.args["encoder"])
        decoder = TransformerDecoder(dtype=self.dtype, **self.args["decoder"])
        self.feature_len = self.args["encoder"]["input_size"]
        self.use_generate = self.args["train_dataset"]["use_generated_data"]
        if not self.use_generate:
            cmvn_file = self.args["normalizer"]["cmvn"]
            mean, istd = load_cmvn(cmvn_file, True)
            mean = torch.from_numpy(mean).float()
            istd = torch.from_numpy(istd).float()
        else:
            mean, istd = None, None
        normalizer = GlobalMVN(self.use_generate, self.feature_len, mean=mean, inv_std=istd)
        loss_fn = LabelSmoothingLoss(**self.args["loss_fn"])
        self.model = Conformer(
            normalizer=normalizer, encoder=encoder, decoder=decoder, loss_fn=loss_fn, args=self.args, dtype=self.dtype
        )
        self.init_type = self.args["trainer"]["init_type"]
        initialize(self.model, self.init_type)
        if self.args["trainer"]["dtype"] == "FLOAT16":
            self.model.half()

    def wrap_pipeline_model(self):
        assert self.model
        self.model.set_start_point_list(self.args["ipu_options"]["pipeline"])

    def build_trainer(self):

        self.trainer = Trainer(
            optimizer=self.optimizer,
            scheduler=self.scheduler,
            model=self.model,
            train_iterator=self.train_iterator,
            val_iterator=self.val_iterator,
            ipu_options=self.ipu_options,
            wandb=self.wandb,
            logger=self.logger,
            args=self.args,
            checkpoint=self.checkpoint,
        )

    def print_args(self):
        pprint.pprint(self.args)

    def train(self):
        self.print_args()
        self.set_random_seed()
        self.launch_popdist()
        self.build_wandb()
        self.build_logger()
        self.build_checkpoints()
        self.build_model()
        self.wrap_pipeline_model()
        self.build_optimizer()
        self.build_iterator()
        self.build_trainer()
        self.trainer.run()

    def validate(self):
        self.print_args()
        self.launch_popdist()
        self.build_wandb()
        self.build_logger()
        self.build_checkpoints()
        self.build_model()
        self.wrap_pipeline_model()
        self.build_optimizer()
        self.build_iterator(False)
        self.build_trainer()
        self.trainer.validate()

    def recognize(self):
        self.print_args()
        self.build_logger()
        self.build_model()
        self.build_iterator(False)
        self.build_trainer()
        self.trainer.recognize()


if __name__ == "__main__":
    fire.Fire(Workflow)