RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment.

[natwille1]

Similar issue to previous (closed) ones but none of those answers worked in my case.

Reproducible example:

resnet = models.resnet50(pretrained=True)
model = SelfSupervisedLearner(
    resnet,
    image_size=IMAGE_SIZE,
    hidden_layer="avgpool",
    projection_size=256,
    projection_hidden_size=4096,
    moving_average_decay=0.99,
)
trainer = pl.Trainer(
    devices=1,
    accelerator="gpu" if torch.cuda.is_available() else "cpu",
    max_epochs=1,
    accumulate_grad_batches=1,
    sync_batchnorm=False,
)
trainer.fit(model, data_loader)
trainer.save_checkpoint("model.pt")
learned_model = SelfSupervisedLearner(resnet, image_size=IMAGE_SIZE).load_from_checkpoint("model.pt")

Full error trace points to this as the root cause:

{
"name": "RuntimeError",
"message": "Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment. If you were attempting to deepcopy a module, this may be because of a torch.nn.utils.weight_norm usage, see pytorch/pytorch#103001",
"stack": "---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[35], line 2
1 #%%
----> 2 learned_model = SelfSupervisedLearner(resnet, image_size=IMAGE_SIZE).load_from_checkpoint("model.pt")
4 # resnet.load_state_dict(torch.load("model.pt")["state_dict"], strict=False)

Cell In[15], line 15, in SelfSupervisedLearner.init(self, net, **kwargs)
13 def init(self, net, **kwargs):
14 super().init()
---> 15 self.learner = BYOL(net, **kwargs)

File ~/anaconda3/envs/google/lib/python3.11/site-packages/byol_pytorch/byol_pytorch.py:235, in BYOL.init(self, net, image_size, hidden_layer, projection_size, projection_hidden_size, augment_fn, augment_fn2, moving_average_decay, use_momentum, sync_batchnorm)
232 self.to(device)
234 # send a mock image tensor to instantiate singleton parameters
--> 235 self.forward(torch.randn(2, 3, image_size, image_size, device=device))

File ~/anaconda3/envs/google/lib/python3.11/site-packages/byol_pytorch/byol_pytorch.py:273, in BYOL.forward(self, x, return_embedding, return_projection)
270 online_pred_one, online_pred_two = online_predictions.chunk(2, dim = 0)
272 with torch.no_grad():
--> 273 target_encoder = self._get_target_encoder() if self.use_momentum else self.online_encoder
275 target_projections, _ = target_encoder(images)
276 target_projections = target_projections.detach()

File ~/anaconda3/envs/google/lib/python3.11/site-packages/byol_pytorch/byol_pytorch.py:28, in singleton..inner_fn..wrapper(self, *args, **kwargs)
25 if instance is not None:
26 return instance
---> 28 instance = fn(self, *args, **kwargs)
29 setattr(self, cache_key, instance)
30 return instance

File ~/anaconda3/envs/google/lib/python3.11/site-packages/byol_pytorch/byol_pytorch.py:239, in BYOL._get_target_encoder(self)
237 @singleton('target_encoder')
238 def _get_target_encoder(self):
--> 239 target_encoder = copy.deepcopy(self.online_encoder)
240 set_requires_grad(target_encoder, False)
241 return target_encoder

File ~/anaconda3/envs/google/lib/python3.11/copy.py:172, in deepcopy(x, memo, _nil)
170 y = x
171 else:
--> 172 y = _reconstruct(x, memo, *rv)
174 # If is its own copy, don't memoize.
175 if y is not x:

File ~/anaconda3/envs/google/lib/python3.11/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
269 if state is not None:
270 if deep:
--> 271 state = deepcopy(state, memo)
272 if hasattr(y, 'setstate'):
273 y.setstate(state)

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
229 memo[id(x)] = y
230 for key, value in x.items():
--> 231 y[deepcopy(key, memo)] = deepcopy(value, memo)
232 return y

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
229 memo[id(x)] = y
230 for key, value in x.items():
--> 231 y[deepcopy(key, memo)] = deepcopy(value, memo)
232 return y

File ~/anaconda3/envs/google/lib/python3.11/copy.py:172, in deepcopy(x, memo, _nil)
170 y = x
171 else:
--> 172 y = _reconstruct(x, memo, *rv)
174 # If is its own copy, don't memoize.
175 if y is not x:

File ~/anaconda3/envs/google/lib/python3.11/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
269 if state is not None:
270 if deep:
--> 271 state = deepcopy(state, memo)
272 if hasattr(y, 'setstate'):
273 y.setstate(state)

[... skipping similar frames: _deepcopy_dict at line 231 (2 times), deepcopy at line 146 (2 times), deepcopy at line 172 (1 times)]

File ~/anaconda3/envs/google/lib/python3.11/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
269 if state is not None:
270 if deep:
--> 271 state = deepcopy(state, memo)
272 if hasattr(y, 'setstate'):
273 y.setstate(state)

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
229 memo[id(x)] = y
230 for key, value in x.items():
--> 231 y[deepcopy(key, memo)] = deepcopy(value, memo)
232 return y

File ~/anaconda3/envs/google/lib/python3.11/copy.py:172, in deepcopy(x, memo, _nil)
170 y = x
171 else:
--> 172 y = _reconstruct(x, memo, *rv)
174 # If is its own copy, don't memoize.
175 if y is not x:

File ~/anaconda3/envs/google/lib/python3.11/copy.py:297, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
295 for key, value in dictiter:
296 key = deepcopy(key, memo)
--> 297 value = deepcopy(value, memo)
298 y[key] = value
299 else:

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:238, in _deepcopy_method(x, memo)
237 def _deepcopy_method(x, memo): # Copy instance methods
--> 238 return type(x)(x.func, deepcopy(x.self, memo))

File ~/anaconda3/envs/google/lib/python3.11/copy.py:172, in deepcopy(x, memo, _nil)
170 y = x
171 else:
--> 172 y = _reconstruct(x, memo, *rv)
174 # If is its own copy, don't memoize.
175 if y is not x:

File ~/anaconda3/envs/google/lib/python3.11/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
269 if state is not None:
270 if deep:
--> 271 state = deepcopy(state, memo)
272 if hasattr(y, 'setstate'):
273 y.setstate(state)

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
229 memo[id(x)] = y
230 for key, value in x.items():
--> 231 y[deepcopy(key, memo)] = deepcopy(value, memo)
232 return y

File ~/anaconda3/envs/google/lib/python3.11/copy.py:146, in deepcopy(x, memo, _nil)
144 copier = _deepcopy_dispatch.get(cls)
145 if copier is not None:
--> 146 y = copier(x, memo)
147 else:
148 if issubclass(cls, type):

File ~/anaconda3/envs/google/lib/python3.11/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
229 memo[id(x)] = y
230 for key, value in x.items():
--> 231 y[deepcopy(key, memo)] = deepcopy(value, memo)
232 return y

File ~/anaconda3/envs/google/lib/python3.11/copy.py:153, in deepcopy(x, memo, _nil)
151 copier = getattr(x, "deepcopy", None)
152 if copier is not None:
--> 153 y = copier(memo)
154 else:
155 reductor = dispatch_table.get(cls)

File ~/anaconda3/envs/google/lib/python3.11/site-packages/torch/_tensor.py:85, in Tensor.deepcopy(self, memo)
83 return handle_torch_function(Tensor.deepcopy, (self,), self, memo)
84 if not self.is_leaf:
---> 85 raise RuntimeError(
86 "Only Tensors created explicitly by the user "
87 "(graph leaves) support the deepcopy protocol at the moment. "
88 "If you were attempting to deepcopy a module, this may be because "
89 "of a torch.nn.utils.weight_norm usage, "
90 "see https://github.com/pytorch/pytorch/pull/103001\"
91 )
92 if id(self) in memo:
93 return memo[id(self)]

RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment. If you were attempting to deepcopy a module, this may be because of a torch.nn.utils.weight_norm usage, see pytorch/pytorch#103001"
}

Please suggest an alternative to deepcopy? Thanks!