Shoud I update trainer in ex. unconditional training

[wasyleque]

Hi,

Im just playing with bunch of images and ex. unconditional training Im saving model and loading it to resume... but ...
Should I also update the trainer ? How often if so ?
Is it unet number 1 in my case ?
trainer.update(unet_number = 1) ???

The original script:

from imagen_pytorch import Unet, Imagen, ImagenTrainer
from imagen_pytorch.data import Dataset

unets for unconditional imagen

unet = Unet(
dim = 32,
dim_mults = (1, 2, 4, 8),
num_resnet_blocks = 1,
layer_attns = (False, False, False, True),
layer_cross_attns = False
)

imagen, which contains the unet above

imagen = Imagen(
condition_on_text = False, # this must be set to False for unconditional Imagen
unets = unet,
image_sizes = 128,
timesteps = 1000
)

trainer = ImagenTrainer(
imagen = imagen,
split_valid_from_train = True # whether to split the validation dataset from the training
).cuda()

instantiate your dataloader, which returns the necessary inputs to the DDPM as tuple in the order of images, text embeddings, then text masks. in this case, only images is returned as it is unconditional training

dataset = Dataset('/path/to/training/images', image_size = 128)

trainer.add_train_dataset(dataset, batch_size = 16)

working training loop

trainer.load('./checkpoint_new.pt') //---> I load the trainer here if I interrupt
for i in range(200000):
loss = trainer.train_step(unet_number = 1, max_batch_size = 4)
print(f'loss: {loss}')

if not (i % 50):
    valid_loss = trainer.valid_step(unet_number = 1, max_batch_size = 4)
    print(f'valid loss: {valid_loss}')

if not (i % 100) and trainer.is_main: # is_main makes sure this can run in distributed
    images = trainer.sample(batch_size = 1, return_pil_images = True) # returns List[Image]
    images[0].save(f'./sample-{i // 100}.png')
    trainer.save('./checkpoint_new.pt') //-------> Saving model each 100 iterations

Second question, script generrates samples automatically, but after leaving training I can not generate images using cmd.
ValueError: unknown imagen type None