malarinv
/
tacotron2
mirror of https://github.com/malarinv/tacotron2
1
0
Fork 0
Code Issues Releases Wiki Activity

single variable for single gpu model

single-gpu-and-0.4
Raul Puri 2018-05-03 22:12:18 -07:00
parent 646ab0d8c8
commit 0d4218fb20
1 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
train.py
View File

@ -128,8 +128,13 @@ def validate(model, criterion, valset, iteration, batch_size, n_gpus,
pin_memory=False, collate_fn=collate_fn) pin_memory=False, collate_fn=collate_fn)
val_loss = 0.0 val_loss = 0.0
if distributed_run or torch.cuda.device_count() > 1:
batch_parser = model.module.parse_batch
else:
batch_parser = model.parse_batch
for i, batch in enumerate(val_loader): for i, batch in enumerate(val_loader):
x, y = model.parse_batch(batch) x, y = batch_parser(batch)
y_pred = model(x) y_pred = model(x)
loss = criterion(y_pred, y) loss = criterion(y_pred, y)
reduced_val_loss = reduce_tensor(loss.data, n_gpus)[0] \ reduced_val_loss = reduce_tensor(loss.data, n_gpus)[0] \
@ -157,6 +162,7 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
if hparams.distributed_run: if hparams.distributed_run:
init_distributed(hparams, n_gpus, rank, group_name) init_distributed(hparams, n_gpus, rank, group_name)
torch.manual_seed(hparams.seed) torch.manual_seed(hparams.seed)
torch.cuda.manual_seed(hparams.seed) torch.cuda.manual_seed(hparams.seed)
@ -188,6 +194,10 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
epoch_offset = max(0, int(iteration / len(train_loader))) epoch_offset = max(0, int(iteration / len(train_loader)))
model.train() model.train()
if distributed_run or torch.cuda.device_count() > 1:
batch_parser = model.module.parse_batch
else:
batch_parser = model.parse_batch
# ================ MAIN TRAINNIG LOOP! =================== # ================ MAIN TRAINNIG LOOP! ===================
for epoch in range(epoch_offset, hparams.epochs): for epoch in range(epoch_offset, hparams.epochs):
print("Epoch: {}".format(epoch)) print("Epoch: {}".format(epoch))
@ -197,7 +207,7 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
param_group['lr'] = learning_rate param_group['lr'] = learning_rate
model.zero_grad() model.zero_grad()
x, y = model.parse_batch(batch) x, y = batch_parser(batch)
y_pred = model(x) y_pred = model(x)
loss = criterion(y_pred, y) loss = criterion(y_pred, y)
reduced_loss = reduce_tensor(loss.data, n_gpus)[0] \ reduced_loss = reduce_tensor(loss.data, n_gpus)[0] \