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

[Fix] Move log value to cpu. #4592

Merged
merged 19 commits into from
Nov 10, 2020
Merged

[Fix] Move log value to cpu. #4592

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
819f213
move value to cpu to save memory
tchaton Nov 9, 2020
2386c0f
update
tchaton Nov 9, 2020
477cd36
move to cpu
tchaton Nov 9, 2020
abbd672
try something
tchaton Nov 9, 2020
fc44142
update
tchaton Nov 9, 2020
e2fb444
Merge branch 'master' into bugfix/4556_gpu_memory_log
tchaton Nov 10, 2020
858ea60
update
tchaton Nov 10, 2020
30afb93
add back out_dict.update({k: v})
tchaton Nov 10, 2020
d9caa5d
add move_metrics_to_cpu
tchaton Nov 10, 2020
3a6244b
Merge branch 'master' into bugfix/4556_gpu_memory_log
tchaton Nov 10, 2020
a57dcfe
update
tchaton Nov 10, 2020
f30bba2
Merge branch 'bugfix/4556_gpu_memory_log' of https://github.com/PyTor…
tchaton Nov 10, 2020
bcd7908
Merge branch 'master' into bugfix/4556_gpu_memory_log
tchaton Nov 10, 2020
eff18c2
Merge branch 'master' into bugfix/4556_gpu_memory_log
tchaton Nov 10, 2020
be4e0f3
Update pytorch_lightning/utilities/memory.py
tchaton Nov 10, 2020
d41273b
resolve comments
tchaton Nov 10, 2020
873e9c0
Update pytorch_lightning/core/step_result.py
tchaton Nov 10, 2020
5e80eb9
Update pytorch_lightning/trainer/connectors/logger_connector/logger_c…
tchaton Nov 10, 2020
5be63ce
Merge branch 'master' into bugfix/4556_gpu_memory_log
SeanNaren Nov 10, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pytorch_lightning/core/step_result.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,10 @@ def log(
if sync_dist and isinstance(value, (torch.Tensor, numbers.Number)):
value = sync_fn(value, group=sync_dist_group, reduce_op=sync_dist_op)

# no need to keep on gpu
if isinstance(value, torch.Tensor) and value.is_cuda:
value = value.cpu()
Copy link
Contributor

Choose a reason for hiding this comment

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

Shouldn't there also be a detach()?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

        # no metrics should be logged with graphs
        if not enable_graph and isinstance(value, torch.Tensor):
            value = value.detach()

        # sync across workers when using distributed training
        sync_fn = sync_fn or sync_ddp_if_available
        if sync_dist and isinstance(value, (torch.Tensor, numbers.Number)):
            value = sync_fn(value, group=sync_dist_group, reduce_op=sync_dist_op)

        # no need to keep on gpu
        if isinstance(value, torch.Tensor) and value.is_cuda:
            value = value.cpu()

detach is called just before.


if 'meta' not in self:
self.__setitem__('meta', {})

Expand Down Expand Up @@ -395,6 +399,11 @@ def detach(self):
if isinstance(v, torch.Tensor):
self.__setitem__(k, v.detach())

def cpu(self):
for k, v in self.items():
tchaton marked this conversation as resolved.
Show resolved Hide resolved
if isinstance(v, torch.Tensor):
self.__setitem__(k, v.cpu())

def __repr__(self):
self_copy = self.copy()

Expand Down
3 changes: 3 additions & 0 deletions pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,6 +392,9 @@ def cache_result(self) -> None:
# attach capture batch_size
Result.attach_batch_size(self._batch_size, hook_result)

hook_result.detach()
hook_result.cpu()
tchaton marked this conversation as resolved.
Show resolved Hide resolved

self._internals[fx_name].append(
hook_result,
dataloader_idx=dataloader_idx,
Expand Down
1 change: 1 addition & 0 deletions pytorch_lightning/trainer/training_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,7 @@ def _process_training_step_output_1_0(self, training_step_output, split_batch):
# track metrics without grads for epoch reduction
training_step_output_for_epoch_end = copy(result)
training_step_output_for_epoch_end.detach()
training_step_output_for_epoch_end.cpu()
tchaton marked this conversation as resolved.
Show resolved Hide resolved

# what flows back into the system
training_step_output = result
Expand Down
1 change: 1 addition & 0 deletions pytorch_lightning/utilities/memory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ def recursive_detach(in_dict: dict) -> dict:
if isinstance(v, dict):
out_dict.update({k: recursive_detach(v)})
elif callable(getattr(v, 'detach', None)):
# detach
out_dict.update({k: v.detach()})
else:
out_dict.update({k: v})
Expand Down