mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 07:41:03 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -0,0 +1,188 @@
|
||||
# *****************************************************************************
|
||||
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
# * Redistributions of source code must retain the above copyright
|
||||
# notice, this list of conditions and the following disclaimer.
|
||||
# * Redistributions in binary form must reproduce the above copyright
|
||||
# notice, this list of conditions and the following disclaimer in the
|
||||
# documentation and/or other materials provided with the distribution.
|
||||
# * Neither the name of the NVIDIA CORPORATION nor the
|
||||
# names of its contributors may be used to endorse or promote products
|
||||
# derived from this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
|
||||
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
||||
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
|
||||
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
|
||||
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
|
||||
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
|
||||
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
||||
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# *****************************************************************************
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import torch
|
||||
|
||||
#=====START: ADDED FOR DISTRIBUTED======
|
||||
from distributed import init_distributed, apply_gradient_allreduce, reduce_tensor
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
#=====END: ADDED FOR DISTRIBUTED======
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
from glow import WaveGlow, WaveGlowLoss
|
||||
from mel2samp import Mel2Samp
|
||||
|
||||
def load_checkpoint(checkpoint_path, model, optimizer):
|
||||
assert os.path.isfile(checkpoint_path)
|
||||
checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
|
||||
iteration = checkpoint_dict['iteration']
|
||||
optimizer.load_state_dict(checkpoint_dict['optimizer'])
|
||||
model_for_loading = checkpoint_dict['model']
|
||||
model.load_state_dict(model_for_loading.state_dict())
|
||||
print("Loaded checkpoint '{}' (iteration {})" .format(
|
||||
checkpoint_path, iteration))
|
||||
return model, optimizer, iteration
|
||||
|
||||
def save_checkpoint(model, optimizer, learning_rate, iteration, filepath):
|
||||
print("Saving model and optimizer state at iteration {} to {}".format(
|
||||
iteration, filepath))
|
||||
model_for_saving = WaveGlow(**waveglow_config).cuda()
|
||||
model_for_saving.load_state_dict(model.state_dict())
|
||||
torch.save({'model': model_for_saving,
|
||||
'iteration': iteration,
|
||||
'optimizer': optimizer.state_dict(),
|
||||
'learning_rate': learning_rate}, filepath)
|
||||
|
||||
def train(num_gpus, rank, group_name, output_directory, epochs, learning_rate,
|
||||
sigma, iters_per_checkpoint, batch_size, seed, fp16_run,
|
||||
checkpoint_path, with_tensorboard):
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
#=====START: ADDED FOR DISTRIBUTED======
|
||||
if num_gpus > 1:
|
||||
init_distributed(rank, num_gpus, group_name, **dist_config)
|
||||
#=====END: ADDED FOR DISTRIBUTED======
|
||||
|
||||
criterion = WaveGlowLoss(sigma)
|
||||
model = WaveGlow(**waveglow_config).cuda()
|
||||
|
||||
#=====START: ADDED FOR DISTRIBUTED======
|
||||
if num_gpus > 1:
|
||||
model = apply_gradient_allreduce(model)
|
||||
#=====END: ADDED FOR DISTRIBUTED======
|
||||
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
||||
|
||||
if fp16_run:
|
||||
from apex import amp
|
||||
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
|
||||
|
||||
# Load checkpoint if one exists
|
||||
iteration = 0
|
||||
if checkpoint_path != "":
|
||||
model, optimizer, iteration = load_checkpoint(checkpoint_path, model,
|
||||
optimizer)
|
||||
iteration += 1 # next iteration is iteration + 1
|
||||
|
||||
trainset = Mel2Samp(**data_config)
|
||||
# =====START: ADDED FOR DISTRIBUTED======
|
||||
train_sampler = DistributedSampler(trainset) if num_gpus > 1 else None
|
||||
# =====END: ADDED FOR DISTRIBUTED======
|
||||
train_loader = DataLoader(trainset, num_workers=1, shuffle=False,
|
||||
sampler=train_sampler,
|
||||
batch_size=batch_size,
|
||||
pin_memory=False,
|
||||
drop_last=True)
|
||||
|
||||
# Get shared output_directory ready
|
||||
if rank == 0:
|
||||
if not os.path.isdir(output_directory):
|
||||
os.makedirs(output_directory)
|
||||
os.chmod(output_directory, 0o775)
|
||||
print("output directory", output_directory)
|
||||
|
||||
if with_tensorboard and rank == 0:
|
||||
from tensorboardX import SummaryWriter
|
||||
logger = SummaryWriter(os.path.join(output_directory, 'logs'))
|
||||
|
||||
model.train()
|
||||
epoch_offset = max(0, int(iteration / len(train_loader)))
|
||||
# ================ MAIN TRAINNIG LOOP! ===================
|
||||
for epoch in range(epoch_offset, epochs):
|
||||
print("Epoch: {}".format(epoch))
|
||||
for i, batch in enumerate(train_loader):
|
||||
model.zero_grad()
|
||||
|
||||
mel, audio = batch
|
||||
mel = torch.autograd.Variable(mel.cuda())
|
||||
audio = torch.autograd.Variable(audio.cuda())
|
||||
outputs = model((mel, audio))
|
||||
|
||||
loss = criterion(outputs)
|
||||
if num_gpus > 1:
|
||||
reduced_loss = reduce_tensor(loss.data, num_gpus).item()
|
||||
else:
|
||||
reduced_loss = loss.item()
|
||||
|
||||
if fp16_run:
|
||||
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
else:
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
print("{}:\t{:.9f}".format(iteration, reduced_loss))
|
||||
if with_tensorboard and rank == 0:
|
||||
logger.add_scalar('training_loss', reduced_loss, i + len(train_loader) * epoch)
|
||||
|
||||
if (iteration % iters_per_checkpoint == 0):
|
||||
if rank == 0:
|
||||
checkpoint_path = "{}/waveglow_{}".format(
|
||||
output_directory, iteration)
|
||||
save_checkpoint(model, optimizer, learning_rate, iteration,
|
||||
checkpoint_path)
|
||||
|
||||
iteration += 1
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-c', '--config', type=str,
|
||||
help='JSON file for configuration')
|
||||
parser.add_argument('-r', '--rank', type=int, default=0,
|
||||
help='rank of process for distributed')
|
||||
parser.add_argument('-g', '--group_name', type=str, default='',
|
||||
help='name of group for distributed')
|
||||
args = parser.parse_args()
|
||||
|
||||
# Parse configs. Globals nicer in this case
|
||||
with open(args.config) as f:
|
||||
data = f.read()
|
||||
config = json.loads(data)
|
||||
train_config = config["train_config"]
|
||||
global data_config
|
||||
data_config = config["data_config"]
|
||||
global dist_config
|
||||
dist_config = config["dist_config"]
|
||||
global waveglow_config
|
||||
waveglow_config = config["waveglow_config"]
|
||||
|
||||
num_gpus = torch.cuda.device_count()
|
||||
if num_gpus > 1:
|
||||
if args.group_name == '':
|
||||
print("WARNING: Multiple GPUs detected but no distributed group set")
|
||||
print("Only running 1 GPU. Use distributed.py for multiple GPUs")
|
||||
num_gpus = 1
|
||||
|
||||
if num_gpus == 1 and args.rank != 0:
|
||||
raise Exception("Doing single GPU training on rank > 0")
|
||||
|
||||
torch.backends.cudnn.enabled = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
train(num_gpus, args.rank, args.group_name, **train_config)
|
||||
Reference in New Issue
Block a user