Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Only put tensors on a device #5223

Merged
merged 2 commits into from Jun 23, 2020
Merged

Only put tensors on a device #5223

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1a06977
Only put tensors on a device
sgugger Jun 23, 2020
10ff478
Type hint and unpack list comprehension
sgugger Jun 23, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 2 additions & 4 deletions src/transformers/trainer.py
View file
Open in desktop
Expand Up @@ -573,8 +573,7 @@ def _training_step(
self, model: nn.Module, inputs: Dict[str, torch.Tensor], optimizer: torch.optim.Optimizer
) -> float:
model.train()
for k, v in inputs.items():
inputs[k] = v.to(self.args.device)
inputs = {k: v.to(self.args.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

Copy link
Member

@julien-c julien-c Jun 23, 2020 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I find this personally slightly hard to read.

Maybe:

for k, v in inputs.items():
    if isinstance(v, torch.Tensor):
        inputs[k] = v.to(self.args.device)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what do you think?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(otherwise, LGTM – we might want to document this in the function's type signature – or not)

Copy link
Collaborator Author

@sgugger sgugger Jun 23, 2020 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Usually using list-comprehensions are faster, but the batches don't have a lot of keys so it probably doesn't matter.

outputs = model(**inputs)
loss = outputs[0] # model outputs are always tuple in transformers (see doc)
Expand Down Expand Up @@ -757,8 +756,7 @@ def _prediction_loop(
for inputs in tqdm(dataloader, desc=description):
has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])

for k, v in inputs.items():
inputs[k] = v.to(self.args.device)
inputs = {k: v.to(self.args.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

with torch.no_grad():
outputs = model(**inputs)
Expand Down