AI4HealthUOL / SSSD

Repository for the paper: 'Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models'
MIT License
270 stars 47 forks source link

Memory location error when calling imputer.impute for CSDIS4Imputer #5

Closed matty6409 closed 1 year ago

matty6409 commented 1 year ago

Hi, ran into the following issue after running imputer.train and imputer.load_weights for the CSDIS4Imputer. Running imputer.impute results in the following error message. The error is thrown when running CSDIS4.py line 2148 when loading the model after training.

RuntimeError: Error(s) in loading state_dict for CSDI_Custom: While copying the parameter named "diffmodel.residual_layers.0.time_layer.s4_layer.kernel.kernel.B", whose dimensions in the model are torch.Size([1, 64, 32, 2]) and whose dimensions in the checkpoint are torch.Size([1, 64, 32, 2]), an exception occurred : ('unsupported operation: more than one element of the written-to tensor refers to a single memory location. Please clone() the tensor before performing the operation.',).

Would love to know a fix for the above issue. Thanks

eruisyan commented 1 year ago

May I ask if you encountered a problem in installing entension cauchy where i stumped , or it goes smoothly when setting up 'entensions cauchy'

image
juanlopezcode commented 1 year ago

Hello @febmatt008, I could speculate that this is related to the S4_lmax hyperparameter. One must adjust this hyperparameter during each training to the length of your sample, otherwise, the S4_lmax will update during training, which will result in an error when trying to load the model weights due to model size missmatch. Regards

jzenn commented 1 year ago

Hey @febmatt008, I just came across the same error you mentioned. The error can be fixed by adapting the register method of SSKernelNPLR as follows.

    def register(self, name, tensor, trainable=False, lr=None, wd=None):
        """Utility method: register a tensor as a buffer or trainable parameter"""

        if trainable:
            self.register_parameter(name, nn.Parameter(tensor.clone()))
        else:
            self.register_buffer(name, tensor.clone())

        optim = {}
        if trainable and lr is not None:
            optim["lr"] = lr
        if trainable and wd is not None:
            optim["weight_decay"] = wd
        if len(optim) > 0:
            setattr(getattr(self, name), "_optim", optim)
kokona1 commented 8 months ago

Hi @jzenn Hi, can you describe in more detail where this method is placed in the code and how it is used, thanks