Open dreamlychina opened 1 year ago
from imagen_pytorch import Unet, Imagen, ImagenTrainer from imagen_pytorch.data import Dataset
output_path="/content/drive/MyDrive/imgen_pytorch/output"
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 = Imagen( condition_on_text = False, # this must be set to False for unconditional Imagen unets = unet, image_sizes = 256, timesteps = 1000 )
trainer = ImagenTrainer( imagen = imagen, split_valid_from_train = True # whether to split the validation dataset from the training ).cuda()
dataset = Dataset('/content/drive/MyDrive/unconditional_generation/dataset_256', image_size = 256)
trainer.add_train_dataset(dataset, batch_size = 16)
for i in range(20000): 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'{output_path}/{i // 100}.png')
This is the training code for your custom dataset .
Thanks for sharing this amazing work,I want to train your inpainting model using my own dataset, could you show me any training script and how to prepare the data at your convenience?