torch-genai 0.0.1.dev1

GenAI

GAN Framework

Generator

Generator must have:

    def forward(self, x, return_q=True):
        """x - images [NCHW]
        return_q - if True, return `r` (recons images [NCHW]),
        and optional `vqres` (quantized info)
        """
        ...
        return r, vqres if return_q else r

VQRes has definition:

VQRes = namedtuple('VQRes', [
    'quantized',
    'indices',
    'entro_mean',
    'mean_entro',
    'entro_loss',
    'commit_loss',
    'cdbook_loss'
])

Discriminator