data_sampler = None
dl = DataLoader(ds, BATCH_SIZE, shuffle=not data_sampler, sampler=data_sampler)
vae_params = dict(image_size=IMAGE_SIZE,
num_layers=NUM_LAYERS,
num_tokens=NUM_TOKENS,
codebook_dim=EMB_DIM,
hidden_dim=HID_DIM,
num_resnet_blocks=NUM_RESNET_BLOCKS)
vae = DiscreteVAE(**vae_params,
smooth_l1_loss=SMOOTH_L1_LOSS,
kl_div_loss_weight=KL_LOSS_WEIGHT)
if not using_deepspeed:
vae = vae.cuda()
assert len(ds) > 0, 'folder does not contain any images'
if distr_backend.is_root_worker():
print(f'{len(ds)} images found for training')
def save_model(path):
if not distr_backend.is_root_worker():
return
save_obj = {'hparams': vae_params, 'weights': vae.state_dict()}
torch.save(save_obj, path)
from torch.utils.data import DataLoader
from dalle_pytorch import DiscreteVAE
imgSize = 256
load_epoch = 280
vae = DiscreteVAE(image_size=imgSize,
num_layers=3,
channels=3,
num_tokens=2048,
codebook_dim=1024,
hidden_dim=128)
vae_dict = torch.load("./models/dvae-" + str(load_epoch) + ".pth")
vae.load_state_dict(vae_dict)
vae.cuda()
batchSize = 12
n_epochs = 500
log_interval = 20
#images = torch.randn(4, 3, 256, 256)
t = transforms.Compose([
transforms.Resize(imgSize),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)) #(0.267, 0.233, 0.234))
])
train_set = datasets.ImageFolder('./imagedata',
# optimizer
get_codes = vae.get_codebook_indices
get_hard_recons = vae.decode
opt = Adam(vae.parameters(), lr = LEARNING_RATE)
sched = ExponentialLR(optimizer = opt, gamma = LR_DECAY_RATE)
logger.debug(f"args.local_rank : {args.local_rank}")
if args.local_rank is not None:
torch.cuda.set_device(args.local_rank)
else:
torch.cuda.set_device(0)
if args.multigpus_distributed:
vae.cuda(args.local_rank)
if args.model_parallel:
vae = smp.DistributedModel(vae)
args.scaler = smp.amp.GradScaler()
opt = smp.DistributedOptimizer(opt)
if args.partial_checkpoint:
args.checkpoint = smp.load(args.partial_checkpoint, partial=True)
vae.load_state_dict(args.checkpoint["model_state_dict"])
opt.load_state_dict(args.checkpoint["optimizer_state_dict"])
elif args.full_checkpoint:
args.checkpoint = smp.load(args.full_checkpoint, partial=False)
vae.load_state_dict(args.checkpoint["model_state_dict"])
opt.load_state_dict(args.checkpoint["optimizer_state_dict"])
else:
