Delete training.py

training.py

@@ -1,254 +0,0 @@
-# ONE EPOCH = one forward pass and one backward pass of all the training examples.
-#
-# BATCH SIZE = the number of training examples in one forward/backward pass. The
-# higher the batch size, the more memory space you'll need.
-#
-# NUMBER OF ITERATIONS = number of passes, each pass using [batch size] number of
-# examples. To be clear, one pass = one forward pass + one backward pass.
-#
-# Example: if you have 1000 training examples, and your batch size is 500, then
-# it will take 2 iterations to complete 1 epoch.
-
-import os
-import time
-import math
-
-import torch
-import torch.distributed as dist
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.data import DataLoader
-from numpy import finfo
-
-from Tacotron2 import tacotron_2
-from fp16_optimizer import FP16_Optimizer
-# from distributed import apply_gradient_allreduce
-from loss_function import Tacotron2Loss
-from logger import Tacotron2Logger
-
-
-def batchnorm_to_float(module):
-    """Converts batch norm modules to FP32"""
-    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
-        module.float()
-    for child in module.children():
-        batchnorm_to_float(child)
-    return module
-
-
-def reduce_tensor(tensor, n_gpus):
-    # this function is recorded in the computation graph. Gradients propagating to the cloned tensor will propagate to
-    # the original tensor
-    rt = tensor.clone()
-    # Each rank has a tensor and all_reduce sums up all tensors from different ranks to all ranks. Computes the average
-    # of the tensor results of all ranks (a rank is a gpu as far as I understood):
-    dist.all_reduce(rt, op=dist.reduce_op.SUM)
-    rt /= n_gpus
-    return rt
-
-
-def prepare_directories_and_logger(output_directory, log_directory, rank):
-    if rank == 0:
-        if not os.path.isdir(output_directory):
-            os.makedirs(output_directory)
-            os.chmod(output_directory, 0o775)
-        logger = Tacotron2Logger(os.path.join(output_directory, log_directory))
-        #  logger = None
-    else:
-        logger = None
-    return logger
-
-
-def warm_start_model(checkpoint_path, model):
-    assert os.path.isfile(checkpoint_path)
-    print("Warm starting model from checkpoint '{}'".format(checkpoint_path))
-    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-    model.load_state_dict(checkpoint_dict['state_dict'])
-    return model
-
-
-def load_checkpoint(checkpoint_path, model, optimizer):
-    assert os.path.isfile(checkpoint_path)
-    print("Loading checkpoint '{}'".format(checkpoint_path))
-    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-    model.load_state_dict(checkpoint_dict['state_dict'])
-    optimizer.load_state_dict(checkpoint_dict['optimizer'])
-    learning_rate = checkpoint_dict['learning_rate']
-    iteration = checkpoint_dict['iteration']
-    print("Loaded checkpoint '{}' from iteration {}".format(checkpoint_path, iteration))
-    return model, optimizer, learning_rate, iteration
-
-
-def save_checkpoint(model, optimizer, learning_rate, iteration, filepath):
-    print("Saving model and optimizer state at iteration {} to {}".format(iteration, filepath))
-    torch.save({'iteration': iteration,
-                'state_dict': model.state_dict(),
-                'optimizer': optimizer.state_dict(),
-                'learning_rate': learning_rate}, filepath)
-
-
-def init_distributed(hyper_params, n_gpus, rank, group_name):
-    assert torch.cuda.is_available(), "Distributed mode requires CUDA"
-    print("Initializing distributed")
-    # Set CUDA device so everything is done on the right GPU
-    torch.cuda.set_device(rank % torch.cuda.device_count())
-
-    # Initialize distributed communication
-    torch.distributed.init_process_group(backend=hyper_params['dist_backend'], rank=rank, world_size=n_gpus,
-                                         init_method=hyper_params['dist_url'], group_name=group_name)
-
-    print("Initializing distributed: Done")
-
-
-def load_model(hyper_params):
-    # according to the documentation, it is recommended to move a model to GPU before constructing the optimizer
-    # model = tacotron_2(hyper_params).cuda()
-    model = tacotron_2(hyper_params)
-    if hyper_params['fp16_run']:  # converts everything into half type (16 bits)
-        model = batchnorm_to_float(model.half())
-        model.decoder.attention_layer.score_mask_value = float(finfo('float16').min)
-
-    # if hyper_params['distributed_run']:
-    #     model = apply_gradient_allreduce(model)
-
-    return model
-
-
-def validate(model, criterion, valset, iteration, batch_size, n_gpus, collate_fn, logger, distributed_run, rank):
-    """Handles all the validation scoring and printing"""
-
-    # We change to eval() because this is an evaluation stage and not a training
-    model.eval()
-    # temporarily set all the requires_grad flag to false
-    with torch.no_grad():
-        # Sampler that restricts data loading to a subset of the dataset. Distributed sampler for distributed batch.
-        # Which samples take (randomization?)
-        val_sampler = DistributedSampler(valset) if distributed_run else None
-        # data loader wraper to the validation data (same as for the training data)
-        val_loader = DataLoader(valset, sampler=val_sampler, num_workers=1, shuffle=False, batch_size=batch_size,
-                                pin_memory=False, collate_fn=collate_fn)
-
-        val_loss = 0.0
-        for i, batch in enumerate(val_loader):
-            x, y = model.parse_batch(batch)
-            y_pred = model(x)
-            _, _, _, _, gst_scores = y_pred
-            if i == 0:
-                validation_gst_scores = gst_scores
-            else:
-                validation_gst_scores = torch.cat((validation_gst_scores, gst_scores), 0)
-            loss = criterion(y_pred, y)
-            if distributed_run:
-                reduced_val_loss = reduce_tensor(loss.data, n_gpus).item()  # gets the pure float value with item()
-            else:
-                reduced_val_loss = loss.item()
-            val_loss += reduced_val_loss
-        val_loss = val_loss / (i + 1)  # Averaged val_loss from all batches
-
-    model.train()
-    if rank == 0:
-        print("Validation loss {}: {:9f} ".format(iteration, val_loss))  # I changed this
-        # print("GST scores of the validation set: {}".format(validation_gst_scores.shape))
-        logger.log_validation(reduced_val_loss, model, y, y_pred, validation_gst_scores, iteration)
-
-
-# ------------------------------------------- MAIN TRAINING METHOD -------------------------------------------------- #
-
-def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus, rank, group_name,
-          hyper_params, train_loader, valset, collate_fn):
-    """Training and validation method with logging results to tensorboard and stdout
-
-    :param output_directory (string): directory to save checkpoints
-    :param log_directory (string): directory to save tensorboard logs
-    :param checkpoint_path (string): checkpoint path
-    :param n_gpus (int): number of gpus
-    :param rank (int): rank of current gpu
-    :param hyper_params (object dictionary): dictionary with all hyper parameters
-    """
-
-    # Check whether is a distributed running
-    if hyper_params['distributed_run']:
-        init_distributed(hyper_params, n_gpus, rank, group_name)
-
-    # set the same fixed seed to reproduce same results everytime we train
-    torch.manual_seed(hyper_params['seed'])
-    torch.cuda.manual_seed(hyper_params['seed'])
-
-    model = load_model(hyper_params)
-    learning_rate = hyper_params['learning_rate']
-    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=hyper_params['weight_decay'])
-
-    if hyper_params['fp16_run']:
-        optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=hyper_params['dynamic_loss_scaling'])
-
-    # Define the criterion of the loss function. The objective.
-    criterion = Tacotron2Loss()
-
-    logger = prepare_directories_and_logger(output_directory, log_directory, rank)
-    # logger = ''
-
-    iteration = 0
-    epoch_offset = 0
-    if checkpoint_path is not None:
-        if warm_start:
-            # Re-start the model from the last checkpoint if we save the parameters and don't want to start from 0
-            model = warm_start_model(checkpoint_path, model)
-        else:
-            # CHECK THIS OUT!!!
-            model, optimizer, _learning_rate, iteration = load_checkpoint(checkpoint_path, model, optimizer)
-            if hyper_params['use_saved_learning_rate']:
-                learning_rate = _learning_rate
-            iteration += 1  # next iteration is iteration + 1
-            epoch_offset = max(0, int(iteration / len(train_loader)))
-
-    # Set this to make all modules and regularization aware this is the training stage:
-    model.train()
-
-    # MAIN LOOP
-    for epoch in range(epoch_offset, hyper_params['epochs']):
-        print("Epoch: {}".format(epoch))
-        for i, batch in enumerate(train_loader):
-            start = time.perf_counter()
-            # CHECK THIS OUT!!!
-            for param_group in optimizer.param_groups:
-                param_group['lr'] = learning_rate
-
-            model.zero_grad()
-            input_data, output_target = model.parse_batch(batch)
-            output_predicted = model(input_data)
-
-            loss = criterion(output_predicted, output_target)
-
-            if hyper_params['distributed_run']:
-                reduced_loss = reduce_tensor(loss.data, n_gpus).item()
-            else:
-                reduced_loss = loss.item()
-
-            if hyper_params['fp16_run']:
-                optimizer.backward(loss)  # transformed optimizer into fp16 type
-                grad_norm = optimizer.clip_fp32_grads(hyper_params['grad_clip_thresh'])
-            else:
-                loss.backward()
-                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), hyper_params['grad_clip_thresh'])
-
-            # Performs a single optimization step (parameter update)
-            optimizer.step()
-            # This boolean controls overflow when running in fp16 optimizer
-            overflow = optimizer.overflow if hyper_params['fp16_run'] else False
-
-            # If overflow is True, it will not enter. If isnan is True, it will not enter neither.
-            if not overflow and not math.isnan(reduced_loss) and rank == 0:
-                duration = time.perf_counter() - start
-                print("Train loss {} {:.6f} Grand Norm {:.6f} {:.2f}s/it".format(iteration, reduced_loss,
-                                                                                 grad_norm, duration))
-                # logs training information of the current iteration
-                logger.log_training(reduced_loss, grad_norm, learning_rate, duration, iteration)
-
-            # Every iters_per_checkpoint steps there is a validation of the model and its updated parameters
-            if not overflow and (iteration % hyper_params['iters_per_checkpoint'] == 0):
-                validate(model, criterion, valset, iteration, hyper_params['batch_size'], n_gpus, collate_fn,
-                         logger, hyper_params['distributed_run'], rank)
-                if rank == 0:
-                    checkpoint_path = os.path.join(output_directory, "checkpoint_{}".format(iteration))
-                    save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path)
-
-            iteration += 1