Add gpu pickleable module to torch #2265
Conversation
mergify
bot
requested review from
a team,
zequnyu and
maliesa96
and removed request for
a team
ziyiwu9494
force-pushed
the
move-before-pickle
branch
from
fe138b1 to
2805fbe
Compare
ziyiwu9494
force-pushed
the
move-before-pickle
branch
from
2805fbe to
3654cf0
Compare
Codecov Report
@@ Coverage Diff @@
## master #2265 +/- ##
==========================================
- Coverage 91.20% 91.17% -0.03%
==========================================
Files 199 199
Lines 11627 10996 -631
Branches 1557 1392 -165
==========================================
- Hits 10604 10026 -578
+ Misses 754 703 -51
+ Partials 269 267 -2
Continue to review full report at Codecov.
|
| save_from_device = global_device() | ||
| self.to('cpu') | ||
| state = self.__dict__.copy() | ||
| state['device'] = save_from_device |
There was a problem hiding this comment.
Probably should call this 'garage.global_device' or something, so that it definitely doesn't conflict with any sub-field names.
| system_device = global_device() | ||
| save_from_device = state['device'] | ||
| if save_from_device == system_device: | ||
| module_state_to(state, system_device) |
There was a problem hiding this comment.
I think you can use the other state-dict moving function you wrote here (even though this kinda isn't a state dict).
There was a problem hiding this comment.
Will every movable parameter of a nn.module be a tensor or a dict?
| """ | ||
| system_device = global_device() | ||
| save_from_device = state['device'] | ||
| if save_from_device == system_device: |
There was a problem hiding this comment.
I'm not sure if this if statement is really needed. I think the idea here is to pre-move everything, as an optimization, but I'm not sure if that's actually faster. If you're going to do this, please use timeit.timeit to measure the performance difference.
There was a problem hiding this comment.
The reason I was doing it this way was because I didn't know if you needed to move any of the other attributes in the dict as well as the module itself, so I'm just moving all internal attributes first and then moving the module itself.
There was a problem hiding this comment.
Oh, I see. .to() should know how to handle that already (since that's how the module was moved to the device in the first place). If modules need to move something besides the default behavior, they should override .to() themselves.
A potential fix for #2079, by creating an garage.torch.Module wrapper class which implements pickling with moving to cpu on pickle and back to original device on unpickle if the device exists.