Fix LRScheduler pickling dragging along CUDA optimizer reference - #1150
Fix LRScheduler pickling dragging along CUDA optimizer reference#1150pctablet505 wants to merge 1 commit into
Conversation
|
Thanks for fixing this old issue.
That was indeed the required ingredient to make the change work. At first I was worried that we'd lose important state by deleting the
The test is a bit indirect, as it doesn't actually test saving a net with CUDA and then loading it on CPU and resuming work. That would be difficult to test though and wouldn't run on the GH CI anyway. I checked that it works locally (also ensuring that it fails without the fix) so at least we have that. Since the PR is still in draft mode, I'm not sure if it's ready yet. Please let me know. |
torch's LR schedulers keep a reference to the optimizer they are built from. Since the LRScheduler callback lives in net.callbacks_, which is not covered by cuda_dependent_attributes_, pickling it used to serialize a duplicate of net.optimizer_ including any device-dependent tensors (e.g. Adam's running averages). This defeated the purpose of handling net.optimizer_ as a CUDA-dependent attribute and broke loading CUDA-trained nets on CPU-only machines. lr_scheduler_ is recreated from scratch in on_train_begin, so it is safe to drop it when pickling the callback. Fixes skorch-dev#1096
f3491cb to
117c284
Compare
|
@pctablet505 I'm still unsure what the state of the PR is from your side? Is it ready? If yes, please mark it as "ready for review" and ping me. |
2 participants
Fixes #1096.
net.optimizer_is handled as a CUDA-dependent attribute when pickling a net, so it can be safely loaded on a CPU-only machine. However, when using theLRSchedulercallback,lr_scheduler_.optimizeris the same optimizer object, and this reference lives insidenet.callbacks_, which is not covered by that CUDA-dependent handling. Pickling the net therefore serialized a second, independent copy of the optimizer (including any device tensors, e.g. Adam's running averages) through the ordinary pickle path, causing CUDA-trained nets with anLRSchedulercallback to fail to load on CPU-only machines.lr_scheduler_is recreated fromnet.optimizer_inon_train_beginanyway, so it doesn't need to be part of the pickled state. This adds a small__getstate__toLRSchedulerthat drops it before pickling, following the same pattern already used byEarlyStoppingandProgressBarfor their own transient state.Added a test that reproduces the reference leak with a stateful optimizer (Adam) and confirms it's gone after a pickle round-trip, and that resuming training afterwards still works correctly.