support convtranspose and activation checkpointing

[guopengf]

#410 support ConvTranspose to avoid F.interpolate limitation and activation checkpointing to save memeory.

[marksgraham]

Hi @guopengf

This looks good to me - thanks for adding. Do you think it would be possible to also implement the checkpointing in the VQ-VAE so we can maintain a consistent interface between the two?

[guopengf]

Hi @marksgraham

I think enabling checkpointing in the VQ-VAE is a good point. I implemented it and added corresponding tests.

However, it is worth noting that currently, TorchScript does not support activation checkpointing. It will trigger the following NotSupportedError error:
torch.jit.frontend.NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults: def checkpoint(function, *args, use_reentrant: bool = True, **kwargs): ~~~~~~~ <--- HERE r"""Checkpoint a model or part of the model 'VQVAE.encode' is being compiled since it was called from 'VQVAE.forward' File "/home/pengfeig/code/GenerativeModels/generative/networks/nets/vqvae.py", line 440 def forward(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: quantizations, quantization_losses = self.quantize(self.encode(images)) ~~~~~~~~~~~~~~~~~~ <--- HERE reconstruction = self.decode(quantizations)
So I removed TorchScript test case here.

I think for a similar reason, Walter removed TorchScript test case to support xformers package in AutoencoderKL.

GenerativeModels/tests/test_autoencoderkl.py

Line 146 in 20d517a

# def test_script(self):