wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-11 11:24:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

[SGD] Support multiple input model (#8246)

Browse Source
This commit is contained in:
Xianyang Liu
2020-05-02 16:49:09 -07:00
committed by GitHub
parent c2acb7ffe2
commit eda526c154
2 changed files with 69 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/util/sgd/torch/training_operator.py
+20 -9
View File
@@ -212,8 +212,9 @@ class TrainingOperator:
updating the model.
By default, this method implementation assumes that batches
are in (features, labels) format. If using amp/fp16
training, it will also scale the loss automatically.
are in (\*features, labels) format. So we also support multiple inputs
model. If using amp/fp16 training, it will also scale the loss
automatically.
You can provide custom loss metrics and training operations if you
override this method. If overriding this method, you can access model,
@@ -237,15 +238,18 @@ class TrainingOperator:
calculate averages.
"""
features, target = batch
# unpack features into list to support multiple inputs model
*features, target = batch
# Create non_blocking tensors for distributed training
if self.use_gpu:
features = features.cuda(non_blocking=True)
features = [
feature.cuda(non_blocking=True) for feature in features
]
target = target.cuda(non_blocking=True)
# Compute output.
with self.timers.record("fwd"):
output = self.model(features)
output = self.model(*features)
loss = self.criterion(output, target)
# Compute gradients in a backward pass.
@@ -261,7 +265,7 @@ class TrainingOperator:
with self.timers.record("apply"):
self.optimizer.step()
return {"train_loss": loss.item(), NUM_SAMPLES: features.size(0)}
return {"train_loss": loss.item(), NUM_SAMPLES: features[0].size(0)}
def validate(self, val_iterator, info):
"""Runs one standard validation pass over the val_iterator.
@@ -304,6 +308,10 @@ class TrainingOperator:
You can override this method to provide arbitrary metrics.
Same as ``train_batch``, this method implementation assumes that
batches are in (\*features, labels) format by default. So we also
support multiple inputs model.
Args:
batch: One item of the validation iterator.
batch_info (dict): Contains information per batch from
@@ -317,15 +325,18 @@ class TrainingOperator:
by default, ``validate`` uses "num_samples" to
calculate averages.
"""
features, target = batch
# unpack features into list to support multiple inputs model
*features, target = batch
if self.use_gpu:
features = features.cuda(non_blocking=True)
features = [
feature.cuda(non_blocking=True) for feature in features
]
target = target.cuda(non_blocking=True)
# compute output
with self.timers.record("eval_fwd"):
output = self.model(features)
output = self.model(*features)
loss = self.criterion(output, target)
_, predicted = torch.max(output.data, 1)