Browse Source

train.py: renaming variable to n_gpus

master
rafaelvalle 6 years ago
parent
commit
0ad65cc053
1 changed files with 4 additions and 4 deletions
  1. +4
    -4
      train.py

+ 4
- 4
train.py View File

@ -28,10 +28,10 @@ def batchnorm_to_float(module):
return module return module
def reduce_tensor(tensor, num_gpus):
def reduce_tensor(tensor, n_gpus):
rt = tensor.clone() rt = tensor.clone()
dist.all_reduce(rt, op=dist.reduce_op.SUM) dist.all_reduce(rt, op=dist.reduce_op.SUM)
rt /= num_gpus
rt /= n_gpus
return rt return rt
@ -135,7 +135,7 @@ def validate(model, criterion, valset, iteration, batch_size, n_gpus,
y_pred = model(x) y_pred = model(x)
loss = criterion(y_pred, y) loss = criterion(y_pred, y)
if distributed_run: if distributed_run:
reduced_val_loss = reduce_tensor(loss.data, num_gpus).item()
reduced_val_loss = reduce_tensor(loss.data, n_gpus).item()
else: else:
reduced_val_loss = loss.item() reduced_val_loss = loss.item()
val_loss += reduced_val_loss val_loss += reduced_val_loss
@ -212,7 +212,7 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
loss = criterion(y_pred, y) loss = criterion(y_pred, y)
if hparams.distributed_run: if hparams.distributed_run:
reduced_loss = reduce_tensor(loss.data, num_gpus).item()
reduced_loss = reduce_tensor(loss.data, n_gpus).item()
else: else:
reduced_loss = loss.item() reduced_loss = loss.item()

Loading…
Cancel
Save