Error: Can't pickle weakref objects

[JCocola]

I am working on autoencoders that implement the gumbel softmax trick.

The code for the autoencoder is borrowed from University of Amsterdam's deep learning course.

In particular, the trainer is defined as

trainer = pl.Trainer(default_root_dir=PATH, gpus=1, max_epochs=500, callbacks=[ModelCheckpoint(save_weights_only=True),
                LearningRateMonitor("epoch")])

Everything works fine when using standard encoders and decoders (MLP or CNNs). But I run into the error TypeError: can't pickle weakref objects when using an encoder that implements the gumbel softmax trick.

The encoder code is taken from here:

class DPS(nn.Module):
       def __init__(self,parameters):
              super(DPS,self).__init__()
              
              # %% initialization
              # running paramters
              self.mux_in = parameters['mux in'] #mux_in means: multiplexer_in
              self.mux_out = parameters['mux out'] #mux_out means: multiplexer_out. We chose those names to indicate the amount of samples coming in (mux_in) and the amount coming out after subsampling (mux_out)
              self.temperature = parameters['temperature']
              self.device = parameters['device']
              self.compression = parameters['compression'] # The factor with which the subsample the data
              
              # create gumble noise constructor
              self.gumbel = torch.distributions.gumbel.Gumbel(0,1)
              
              # create the logit parameter
              self.logits = nn.Parameter(torch.randn(self.mux_in,)/self.compression)
              
              
       def forward(self,x):
              # x is the input of the actual input of the entire model. It directly enters this sampling layer.
              # %% forward call
              #find the batch size
              batch_size = x.size(0)
              
              #create perturbed logits
              noise = self.gumbel.sample((batch_size,self.mux_in,1)).repeat(1,1,self.mux_out).to(self.device)
              perturbed_logits = torch.transpose(self.logits.repeat(batch_size,self.mux_out,1),1,2) + noise
              
              #find the locations of the topk samples
              _,topk_indices = torch.topk(perturbed_logits[:,:,0],self.mux_out)
              
              # %% counting vectors (helpers)
              batch_indices = torch.tensor([range(batch_size)]).reshape(batch_size,1).repeat(1,self.mux_out)
              k_counter = torch.tensor([range(self.mux_out)])

              # %% create  hard sample matrix
              hard_sample_matrix = torch.zeros(batch_size,self.mux_in).to(self.device)
              hard_sample_matrix[batch_indices,topk_indices] = 1
              
              # %% create the soft sample matrix
              masking_matrix = torch.zeros(batch_size,self.mux_in,self.mux_out).to(self.device)

              #cumulative masking:
			  # note that the masking in the pytorch implementation is done in the logits space. The tf implementation does this in the probability space.
			  # It should however not make a great difference, it is just a different implementation of the same masking behaviour.
              masking_matrix[batch_indices,topk_indices,k_counter] = -np.inf
              masking_matrix = torch.cumsum(masking_matrix, dim = 2)
              masking_matrix[batch_indices,topk_indices,k_counter] = 0 # mimic the exclusive=True parameter from tf.cumsum
              
              masked_logits = perturbed_logits + masking_matrix
              
              # softmaxing
              soft_sample_matrix = F.softmax(masked_logits/self.temperature,dim = 1)
              
              # collapse
              soft_sample_matrix_collapsed = torch.sum(soft_sample_matrix,dim=2)
              
              # %% Make sure that in the forward pass the hard sample matrix will be passed, whereas in the backwards pass, only the softsample matrix will be seen for the gradient. 
			  # The hard_sample_matrix already has no defined gradient, so doesn't need to be detached explicitly.
              sample_matrix = hard_sample_matrix - soft_sample_matrix_collapsed.detach() + soft_sample_matrix_collapsed 
              
              # %% element-wise multiplication of the input tensor with the binary mask. The non-sampled values are now set at zero.
              y = sample_matrix * x
			  
			  # Note that this layer outputs the already subsampled (and zero-filled) input signal. This is different than the output of the tf implemenation of the topKsampling layer.
			  # That layer outputs a hard (or for back prop a soft) sampling matrix of size MxN. Only after this layer a sampling mask is created from that and element-wise applied on the input data.

              return y

---

Full error message is below

TypeError                                 Traceback (most recent call last)
[<ipython-input-22-dd0e96f27375>](https://localhost:8080/#) in <module>()
      3 start = datetime.now()
      4 print(start)
----> 5 model, result = train_AE(trainer, model, train_loader, val_loader, test_loader)
      6 delta = datetime.now() - start
      7 

7 frames
[<ipython-input-12-ea4480eb4109>](https://localhost:8080/#) in train_AE(trainer, model, train_loader, val_loader, test_loader)
     21 
     22 def train_AE(trainer, model, train_loader, val_loader, test_loader):
---> 23     trainer.fit(model, train_loader, val_loader)
     24 
     25     # Test best model on validation and test set

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    769         self.strategy.model = model
    770         self._call_and_handle_interrupt(
--> 771             self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    772         )
    773 

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    721                 return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
    722             else:
--> 723                 return trainer_fn(*args, **kwargs)
    724         # TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7
    725         except KeyboardInterrupt as exception:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    809             ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
    810         )
--> 811         results = self._run(model, ckpt_path=self.ckpt_path)
    812 
    813         assert self.state.stopped

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py](https://localhost:8080/#) in _run(self, model, ckpt_path)
   1150         # clean hparams
   1151         if hasattr(model, "hparams"):
-> 1152             parsing.clean_namespace(model.hparams)
   1153 
   1154         # attach model to the strategy

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/parsing.py](https://localhost:8080/#) in clean_namespace(hparams)
    102         hparams_dict = hparams.__dict__
    103 
--> 104     del_attrs = [k for k, v in hparams_dict.items() if not is_picklable(v)]
    105 
    106     for k in del_attrs:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/parsing.py](https://localhost:8080/#) in <listcomp>(.0)
    102         hparams_dict = hparams.__dict__
    103 
--> 104     del_attrs = [k for k, v in hparams_dict.items() if not is_picklable(v)]
    105 
    106     for k in del_attrs:

[/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/parsing.py](https://localhost:8080/#) in is_picklable(obj)
     89 
     90     try:
---> 91         pickle.dumps(obj)
     92         return True
     93     except (pickle.PickleError, AttributeError, RuntimeError):

TypeError: can't pickle weakref objects