def generate_images(captions):
print(sum(captions, start=[]))
all_words = list(sorted(frozenset(sum(captions, start=[]))))
word_tokens = dict(zip(all_words, range(1, len(all_words) + 1)))
caption_tokens = [[word_tokens[w] for w in c] for c in captions]
logging.info(f"{all_words =}")
logging.info(f"{word_tokens =}")
logging.info(f"{caption_tokens =}")
longest_caption = max(len(c) for c in captions)
captions_array = np.zeros((len(caption_tokens), longest_caption),
dtype=np.int64)
for i in range(len(caption_tokens)):
captions_array[i, :len(caption_tokens[i])] = caption_tokens[i]
# captions_array = torch.from_numpy(captions_array).cuda()
captions_array = torch.from_numpy(captions_array)
captions_mask = captions_array != 0
logging.info(f"{captions_array = }")
dalle = DALLE(
dim=1024,
vae=
vae, # automatically infer (1) image sequence length and (2) number of image tokens
num_text_tokens=len(word_tokens) + 1, # vocab size for text
text_seq_len=longest_caption, # text sequence length
depth=12, # should aim to be 64
heads=16, # attention heads
dim_head=64, # attention head dimension
attn_dropout=0.1, # attention dropout
ff_dropout=0.1, # feedforward dropout
)
# .cuda
generated_image_codes = []
with torch.no_grad():
for i in range(0, len(captions), 128):
generated = generate_image_code(
dalle,
captions_array[i:i + 128, ...],
mask=captions_mask[i:i + 128, ...],
)
generated_image_codes.append(generated)
generated_image_codes = torch.cat(generated_image_codes, axis=0)
with torch.no_grad():
generated_images = vae.decode(generated_image_codes)
logging.info(f"{generated_images = }")
def get_dalle(vae, vocab, args):
dalle = DALLE(dim=args.codebook_dims,
vae=vae,
num_text_tokens=len(vocab) + 1,
text_seq_len=len(vocab),
depth=16,
heads=8,
dim_head=64,
attn_dropout=0.1,
ff_dropout=0.1,
reversible=True)
if args.dalle is not None and os.path.isfile(args.dalle):
print(f"loading state dict from {args.dalle}")
dalle.load_state_dict(torch.load(args.dalle))
dalle.to(args.device)
vae.to(args.device)
return dalle
data_sampler = torch.utils.data.distributed.DistributedSampler(
ds,
num_replicas=distr_backend.get_world_size(),
rank=distr_backend.get_rank())
else:
data_sampler = None
dl = DataLoader(ds,
batch_size=BATCH_SIZE,
shuffle=not data_sampler,
drop_last=True,
sampler=data_sampler)
# initialize DALL-E
dalle = DALLE(vae=vae, **dalle_params)
if args.fp16:
dalle = dalle.half()
dalle = dalle.cuda()
if RESUME:
dalle.load_state_dict(weights)
# optimizer
opt = Adam(dalle.parameters(), lr=LEARNING_RATE)
if LR_DECAY:
scheduler = ReduceLROnPlateau(
opt,
mode="min",
assert dalle_path.exists(), 'trained DALL-E must exist'
load_obj = torch.load(str(dalle_path))
dalle_params, vae_params, weights = load_obj.pop('hparams'), load_obj.pop('vae_params'), load_obj.pop('weights')
dalle_params.pop('vae', None) # cleanup later
if vae_params is not None:
vae = DiscreteVAE(**vae_params)
elif not args.taming:
vae = OpenAIDiscreteVAE()
else:
vae = VQGanVAE1024()
dalle = DALLE(vae = vae, **dalle_params).cuda()
dalle.load_state_dict(weights)
# generate images
image_size = vae.image_size
texts = args.text.split('|')
for text in tqdm(texts):
text = tokenizer.tokenize([args.text], dalle.text_seq_len).cuda()
text = repeat(text, '() n -> b n', b = args.num_images)
outputs = []
hidden_dim=64, # hidden dimension
num_resnet_blocks=1, # number of resnet blocks
temperature=
0.9, # gumbel softmax temperature, the lower this is, the harder the discretization
straight_through=
False, # straight-through for gumbel softmax. unclear if it is better one way or the other
).cuda()
vae.load_state_dict(torch.load("Vae-small.pth"))
dalle = DALLE(
dim=1024,
vae=
vae, # automatically infer (1) image sequence length and (2) number of image tokens
num_text_tokens=NUM_TOKENS, # vocab size for text
text_seq_len=TEXTSEQLEN, # text sequence length
depth=12, # should aim to be 64
heads=16, # attention heads
dim_head=64, # attention head dimension
attn_dropout=0.1, # attention dropout
ff_dropout=0.1 # feedforward dropout
).cuda()
dalle.load_state_dict(torch.load("dalle-small.pth"))
"""
text = torch.randint(0, NUM_TOKENS, (BATCH_SIZE, TEXTSEQLEN))
images = torch.randn(BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE)
mask = torch.ones_like(text).bool()
"""
tokenDset = token_dataset('./coco/merged-smallsample.txt')
# create dataset and dataloader
ds = TextImageDataset(
args.image_text_folder,
text_len = TEXT_SEQ_LEN,
image_size = IMAGE_SIZE
)
assert len(ds) > 0, 'dataset is empty'
print(f'{len(ds)} image-text pairs found for training')
dl = DataLoader(ds, batch_size = BATCH_SIZE, shuffle = True, drop_last = True)
# initialize DALL-E
dalle = DALLE(**dalle_params).cuda()
if exists(args.dalle_path):
dalle.load_state_dict(weights)
# optimizer
opt = Adam(dalle.parameters(), lr = LEARNING_RATE)
# experiment tracker
import wandb
wandb.config.depth = DEPTH
wandb.config.heads = HEADS
wandb.config.dim_head = DIM_HEAD
EPOCHS * DATASET_SIZE))
loss = vae(img, return_recon_loss=True)
VAEloss.append(loss.cpu().detach().numpy())
loss.backward()
optimizerVAE.step()
np.savetxt("vaeloss.csv", np.asarray(VAEloss), delimiter=",")
torch.save(vae.state_dict(), "Vae-small.pth")
dalle = DALLE(
dim=1024,
vae=
vae, # automatically infer (1) image sequence length and (2) number of image tokens
num_text_tokens=NUM_TOKENS, # vocab size for text
text_seq_len=TEXTSEQLEN, # text sequence length
depth=12, # should aim to be 64
heads=16, # attention heads
dim_head=64, # attention head dimension
attn_dropout=0.1, # attention dropout
ff_dropout=0.1 # feedforward dropout
).cuda()
optimizerDALLE = torch.optim.Adam(dalle.parameters(), lr=learning_rate)
DALLEloss = []
for epoch in range(EPOCHS):
for i in range(DATASET_SIZE):
#print(i,":",tokenDset.getRand(i),img.size())
optimizerDALLE.zero_grad()
img, strs = cap[i]
#print(img.size())
data_sampler = None
if ENABLE_WEBDATASET:
# WebLoader for WebDataset and DeepSpeed compatibility
dl = wds.WebLoader(ds, batch_size=None, shuffle=False) # optionally add num_workers=2 (n) argument
number_of_batches = DATASET_SIZE // (BATCH_SIZE * distr_backend.get_world_size())
dl = dl.repeat(2).slice(number_of_batches)
dl.length = number_of_batches
else:
# Regular DataLoader for image-text-folder datasets
dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=is_shuffle, drop_last=True, sampler=data_sampler)
# initialize DALL-E
dalle = DALLE(vae=vae, **dalle_params)
if not using_deepspeed:
if args.fp16:
dalle = dalle.half()
dalle = dalle.cuda()
if RESUME and not using_deepspeed:
dalle.load_state_dict(weights)
# optimizer
opt = Adam(get_trainable_params(dalle), lr=LEARNING_RATE)
if RESUME and opt_state:
opt.load_state_dict(opt_state)
if LR_DECAY:
def main(args):
# constants
VAE_PATH = args.vae_path
DALLE_PATH = args.dalle_path
RESUME = exists(DALLE_PATH)
EPOCHS = args.epochs
BATCH_SIZE = args.batch_size
LEARNING_RATE = args.learning_rate
GRAD_CLIP_NORM = args.clip_grad_norm
LR_DECAY = args.lr_decay
MODEL_DIM = args.dim
TEXT_SEQ_LEN = args.text_seq_len
DEPTH = args.depth
HEADS = args.heads
DIM_HEAD = args.dim_head
REVERSIBLE = args.reversible
LOSS_IMG_WEIGHT = args.loss_img_weight
ATTN_TYPES = tuple(args.attn_types.split(','))
DEEPSPEED_CP_AUX_FILENAME = 'auxiliary.pt'
# initialize distributed backend
# initialize distributed backend
if args.sagemakermp:
args.deepspeed = False
using_deepspeed = False
else:
args.deepspeed = True
distr_backend = distributed_utils.set_backend_from_args(args)
distr_backend.initialize(args)
if args.sagemakermp:
args = smp_init(args)
distributed_utils.using_backend(distributed_utils.SageMakerMPBackend)
else:
using_deepspeed = \
distributed_utils.using_backend(distributed_utils.DeepSpeedBackend)
args.rank = int(os.environ.get('RANK'))
args.world_size = int(os.environ.get('WORLD_SIZE'))
args.local_rank = int(os.environ.get('LOCAL_RANK'))
args.global_rank = args.rank
logger.debug(f"using_deepspeed : {using_deepspeed}")
logger.debug(
f"args.local_rank : {args.local_rank}, args.rank : {args.rank}")
# tokenizer
logger.debug(f"exists(args.bpe_path) : {exists(args.bpe_path)}, args.chinese : {args.chinese}")
if exists(args.bpe_path):
klass = HugTokenizer if args.hug else YttmTokenizer
tokenizer = klass(args.bpe_path)
elif args.chinese:
tokenizer = ChineseTokenizer()
else:
tokenizer = SimpleTokenizer()
# reconstitute vae
if RESUME:
dalle_path = Path(DALLE_PATH)
if using_deepspeed:
cp_dir = cp_path_to_dir(dalle_path, 'ds')
assert cp_dir.is_dir(), \
f'DeepSpeed checkpoint directory {cp_dir} not found'
dalle_path = cp_dir / DEEPSPEED_CP_AUX_FILENAME
else:
assert dalle_path.exists(), 'DALL-E model file does not exist'
loaded_obj = torch.load(str(dalle_path), map_location='cpu')
dalle_params, vae_params, weights = loaded_obj['hparams'], loaded_obj[
'vae_params'], loaded_obj['weights']
if vae_params is not None:
vae = DiscreteVAE(**vae_params)
else:
vae_klass = OpenAIDiscreteVAE if not args.taming else VQGanVAE1024
vae = vae_klass(args)
dalle_params = dict(**dalle_params)
IMAGE_SIZE = vae.image_size
else:
if exists(VAE_PATH):
vae_path = Path(VAE_PATH)
assert vae_path.exists(), 'VAE model file does not exist'
assert not vae_path.is_dir(), \
('Cannot load VAE model from directory; please use a '
'standard *.pt checkpoint. '
'Currently, merging a DeepSpeed-partitioned VAE into a DALLE '
'model is not supported.')
loaded_obj = torch.load(str(vae_path))
vae_params, weights = loaded_obj['hparams'], loaded_obj['weights']
vae = DiscreteVAE(**vae_params)
vae.load_state_dict(weights)
else:
if args.rank == 0:
# if distr_backend.is_root_worker():
print('using pretrained VAE for encoding images to tokens')
vae_params = None
logger.debug(f"************* args.taming : {args.taming}")
vae_klass = OpenAIDiscreteVAE if not args.taming else VQGanVAE1024
vae = vae_klass(args)
IMAGE_SIZE = vae.image_size
dalle_params = dict(
num_text_tokens=tokenizer.vocab_size,
text_seq_len=TEXT_SEQ_LEN,
dim=MODEL_DIM,
depth=DEPTH,
heads=HEADS,
dim_head=DIM_HEAD,
reversible=REVERSIBLE,
loss_img_weight=LOSS_IMG_WEIGHT,
attn_types=ATTN_TYPES,
)
# configure OpenAI VAE for float16s
if isinstance(vae, OpenAIDiscreteVAE) and args.fp16:
vae.enc.blocks.output.conv.use_float16 = True
# create dataset and dataloader
is_shuffle = not distributed_utils.using_backend(
distributed_utils.HorovodBackend)
ds = TextImageDataset(
args.image_text_folder,
text_len=TEXT_SEQ_LEN,
image_size=IMAGE_SIZE,
resize_ratio=args.resize_ratio,
truncate_captions=args.truncate_captions,
tokenizer=tokenizer,
shuffle=is_shuffle,
)
assert len(ds) > 0, 'dataset is empty'
# if distr_backend.is_root_worker():
if args.rank == 0:
print(f'{len(ds)} image-text pairs found for training')
if not is_shuffle:
data_sampler = torch.utils.data.distributed.DistributedSampler(
ds, num_replicas=args.world_size, rank=args.rank)
elif args.sagemakermp:
args.ds = ds
ds = split_dataset(args)
data_sampler = None
else:
data_sampler = None
print(f"data_sampler : {data_sampler}")
# uncorrectable NVLink error was detected during the execution --> remove
kwargs = {'num_workers': args.num_worker, 'pin_memory': True}
dl = DataLoader(ds,
batch_size=BATCH_SIZE,
shuffle=is_shuffle,
drop_last=True,
sampler=data_sampler,
**kwargs
)
logger.info("Processes {}/{} ({:.0f}%) of train data".format(
len(dl.sampler), len(dl.dataset),
100. * len(dl.sampler) / len(dl.dataset)))
# initialize DALL-E
dalle = DALLE(vae=vae, **dalle_params)
if not using_deepspeed:
if args.fp16:
dalle = dalle.half()
dalle = dalle.cuda()
if RESUME and not using_deepspeed:
dalle.load_state_dict(weights)
# optimizer
opt = Adam(get_trainable_params(dalle), lr=LEARNING_RATE)
if LR_DECAY:
scheduler = ReduceLROnPlateau(
opt,
mode="min",
factor=0.5,
patience=10,
cooldown=10,
min_lr=1e-6,
verbose=True,
)
# if distr_backend.is_root_worker():
if args.global_rank == 0:
# experiment tracker
model_config = dict(depth=DEPTH, heads=HEADS, dim_head=DIM_HEAD)
logger.debug(f"args.wandb_name : {args.wandb_name}, RESUME : {RESUME}")
run = wandb.init(
project=args.wandb_name, # 'dalle_train_transformer' by default
resume=RESUME,
config=model_config,
)
# distribute
distr_backend.check_batch_size(BATCH_SIZE)
deepspeed_config = {
'train_batch_size': BATCH_SIZE,
'gradient_clipping': GRAD_CLIP_NORM,
'fp16': {
'enabled': args.fp16,
},
}
(distr_dalle, distr_opt, distr_dl,
distr_scheduler) = distr_backend.distribute(
args=args,
model=dalle,
optimizer=opt,
model_parameters=get_trainable_params(dalle),
training_data=ds if using_deepspeed else dl,
lr_scheduler=scheduler if LR_DECAY else None,
config_params=deepspeed_config,
)
avoid_model_calls = using_deepspeed and args.fp16
if args.sagemakermp:
args.distr_dalle = smp.DistributedModel(distr_dalle)
args.scaler = smp.amp.GradScaler()
args.distr_opt = smp.DistributedOptimizer(distr_opt)
if RESUME and using_deepspeed:
distr_dalle.load_checkpoint(str(cp_dir))
# training
for epoch in range(EPOCHS):
logger.debug(f"********* epoch : {epoch} **********")
if data_sampler:
data_sampler.set_epoch(epoch)
for i, (text, images) in enumerate(distr_dl):
if args.fp16:
images = images.half()
text, images = map(lambda t: t.cuda(), (text, images))
if args.sagemakermp:
args.distr_opt.zero_grad()
loss = train_step(args, text, images, return_loss=True)
loss = loss.reduce_mean()
else:
loss = distr_dalle(text, images, return_loss=True, args=args)
if using_deepspeed:
distr_dalle.backward(loss)
distr_dalle.step()
# Gradients are automatically zeroed after the step
elif args.sagemakermp:
if args.amp:
scaler.step(args.distr_opt)
scaler.update()
else:
# some optimizers like adadelta from PT 1.8 dont like it when optimizer.step is called with no param
if len(list(args.distr_dalle.local_parameters())) > 0:
args.distr_opt.step()
else:
loss.backward()
clip_grad_norm_(distr_dalle.parameters(), GRAD_CLIP_NORM)
distr_opt.step()
distr_opt.zero_grad()
# Collective loss, averaged
avg_loss = distr_backend.average_all(loss)
log = {}
# if i % 10 == 0 and distr_backend.is_root_worker():
if i % 10 == 0 and args.rank == 0:
print(epoch, i, f'loss - {avg_loss.item()}')
log = {
**log, 'epoch': epoch,
'iter': i,
'loss': avg_loss.item()
}
if i % 100 == 0:
# if distr_backend.is_root_worker():
if args.rank == 0:
sample_text = text[:1]
token_list = sample_text.masked_select(
sample_text != 0).tolist()
decoded_text = tokenizer.decode(token_list)
logger.debug(f"******* avoid_model_calls : {avoid_model_calls}")
if not avoid_model_calls:
# CUDA index errors when we don't guard this
image = dalle.generate_images(
text[:1], filter_thres=0.9) # topk sampling at 0.9
wandb.save(f'./dalle.pt')
log = {
**log,
}
if not avoid_model_calls:
log['image'] = wandb.Image(image, caption=decoded_text)
args.distr_dalle = distr_dalle
args.dalle_params = dalle_params
args.vae_params = vae_params
args.using_deepspeed = using_deepspeed
args.DEEPSPEED_CP_AUX_FILENAME = DEEPSPEED_CP_AUX_FILENAME
save_model(args, f'{args.model_dir}/dalle.pt')
# if distr_backend.is_root_worker():
if args.rank == 0:
wandb.log(log)
# text, images = prefetcher.next()
if LR_DECAY and not using_deepspeed:
# Scheduler is automatically progressed after the step when
# using DeepSpeed.
distr_scheduler.step(loss)
# if distr_backend.is_root_worker():
if args.global_rank == 0:
# save trained model to wandb as an artifact every epoch's end
model_artifact = wandb.Artifact('trained-dalle',
type='model',
metadata=dict(model_config))
# model_artifact.add_file('dalle.pt')
run.log_artifact(model_artifact)
args.distr_dalle = distr_dalle
args.dalle_params = dalle_params
args.vae_params = vae_params
args.using_deepspeed = using_deepspeed
args.DEEPSPEED_CP_AUX_FILENAME = DEEPSPEED_CP_AUX_FILENAME
save_model(args, f'{args.model_dir}/dalle-final.pt')
# if distr_backend.is_root_worker():
if args.global_rank == 0:
wandb.save('./dalle-final.pt')
model_artifact = wandb.Artifact('trained-dalle',
type='model',
metadata=dict(model_config))
# model_artifact.add_file('dalle-final.pt')
run.log_artifact(model_artifact)
wandb.finish()
codebook_dim=256,
hidden_dim=128,
temperature=0.9)
# load pretrained vae
vae_dict = torch.load("./models/" + vaename + "-" + str(load_epoch) + ".pth")
vae.load_state_dict(vae_dict)
vae.to(device)
dalle = DALLE(
dim=256, #512,
vae=
vae, # automatically infer (1) image sequence length and (2) number of image tokens
num_text_tokens=10000, # vocab size for text
text_seq_len=256, # text sequence length
depth=6, # should be 64
heads=8, # attention heads
dim_head=64, # attention head dimension
attn_dropout=0.1, # attention dropout
ff_dropout=0.1 # feedforward dropout
)
# load pretrained dalle if continuing training
dalle_dict = torch.load(loadfn)
dalle.load_state_dict(dalle_dict)
dalle.to(device)
# get image and text data
def main(args):
# constants
print(f"torch.cuda.nccl.version() : {torch.cuda.nccl.version()}")
VAE_PATH = args.vae_path
DALLE_PATH = args.dalle_path
RESUME = exists(DALLE_PATH)
EPOCHS = args.epochs
BATCH_SIZE = args.batch_size
LEARNING_RATE = args.learning_rate
GRAD_CLIP_NORM = args.clip_grad_norm
LR_DECAY = args.lr_decay
SAVE_EVERY_N_STEPS = args.save_every_n_steps
MODEL_DIM = args.dim
TEXT_SEQ_LEN = args.text_seq_len
DEPTH = args.depth
HEADS = args.heads
DIM_HEAD = args.dim_head
REVERSIBLE = args.reversible
LOSS_IMG_WEIGHT = args.loss_img_weight
ATTN_TYPES = tuple(args.attn_types.split(','))
DEEPSPEED_CP_AUX_FILENAME = 'auxiliary.pt'
DALLE_OUTPUT_FILE_NAME = args.dalle_output_file_name
# initialize distributed backend
args.deepspeed = True
distr_backend = distributed_utils.set_backend_from_args(args)
distr_backend.initialize(args)
using_deepspeed = \
distributed_utils.using_backend(distributed_utils.DeepSpeedBackend)
args.rank = int(os.environ.get('RANK'))
args.world_size = int(os.environ.get('WORLD_SIZE'))
args.local_rank = int(os.environ.get('LOCAL_RANK'))
args.global_rank = args.rank
logger.debug(f"using_deepspeed : {using_deepspeed}")
logger.debug(
f"args.local_rank : {args.local_rank}, args.rank : {args.rank}")
print(
f"********* torch.distributed.get_rank() : {torch.distributed.get_rank()}"
)
# tokenizer
logger.debug(
f"exists(args.bpe_path) : {exists(args.bpe_path)}, args.chinese : {args.chinese}"
)
# if args.local_rank == 0 or args.local_rank == 1:
# # print(f"args.job_name : {args.job_name}")
# gpu_mon_thread = aws_util.GPUMon(device_index= args.local_rank, job_name=args.job_name)
# gpu_mon_thread.start()
if exists(args.bpe_path):
klass = HugTokenizer if args.hug else YttmTokenizer
tokenizer = klass(args.bpe_path)
elif args.chinese:
tokenizer = ChineseTokenizer()
else:
tokenizer = SimpleTokenizer()
# reconstitute vae
if RESUME:
dalle_path = Path(DALLE_PATH)
if using_deepspeed:
cp_dir = cp_path_to_dir(dalle_path, 'ds')
assert cp_dir.is_dir(), \
f'DeepSpeed checkpoint directory {cp_dir} not found'
dalle_path = cp_dir / DEEPSPEED_CP_AUX_FILENAME
else:
assert dalle_path.exists(), 'DALL-E model file does not exist'
loaded_obj = torch.load(str(dalle_path), map_location='cpu')
dalle_params, vae_params, weights = loaded_obj['hparams'], loaded_obj[
'vae_params'], loaded_obj['weights']
if vae_params is not None:
vae = DiscreteVAE(**vae_params)
else:
vae_klass = OpenAIDiscreteVAE if not args.taming else VQGanVAE1024
vae = vae_klass(args)
dalle_params = dict(**dalle_params)
IMAGE_SIZE = vae.image_size
else:
if exists(VAE_PATH):
vae_path = Path(VAE_PATH)
assert vae_path.exists(), 'VAE model file does not exist'
assert not vae_path.is_dir(), \
('Cannot load VAE model from directory; please use a '
'standard *.pt checkpoint. '
'Currently, merging a DeepSpeed-partitioned VAE into a DALLE '
'model is not supported.')
loaded_obj = torch.load(str(vae_path))
vae_params, weights = loaded_obj['hparams'], loaded_obj['weights']
vae = DiscreteVAE(**vae_params)
vae.load_state_dict(weights)
else:
if args.rank == 0:
# if distr_backend.is_root_worker():
print('using pretrained VAE for encoding images to tokens')
vae_params = None
logger.debug(f"************* args.taming : {args.taming}")
vae_klass = OpenAIDiscreteVAE if not args.taming else VQGanVAE1024
vae = vae_klass(args)
IMAGE_SIZE = vae.image_size
dalle_params = dict(
num_text_tokens=tokenizer.vocab_size,
text_seq_len=TEXT_SEQ_LEN,
dim=MODEL_DIM,
depth=DEPTH,
heads=HEADS,
dim_head=DIM_HEAD,
reversible=REVERSIBLE,
loss_img_weight=LOSS_IMG_WEIGHT,
attn_types=ATTN_TYPES,
)
# configure OpenAI VAE for float16s
if isinstance(vae, OpenAIDiscreteVAE) and args.fp16:
vae.enc.blocks.output.conv.use_float16 = True
# create dataset and dataloader
is_shuffle = not distributed_utils.using_backend(
distributed_utils.HorovodBackend)
ds = TextImageDataset(
args.image_text_folder,
text_len=TEXT_SEQ_LEN,
image_size=IMAGE_SIZE,
resize_ratio=args.resize_ratio,
truncate_captions=args.truncate_captions,
tokenizer=tokenizer,
shuffle=is_shuffle,
)
assert len(ds) > 0, 'dataset is empty'
# if distr_backend.is_root_worker():
if args.rank == 0:
print(f'{len(ds)} image-text pairs found for training')
if not is_shuffle:
data_sampler = torch.utils.data.distributed.DistributedSampler(
ds, num_replicas=args.world_size, rank=args.rank)
else:
data_sampler = None
kwargs = {'num_workers': args.num_worker, 'pin_memory': True}
dl = DataLoader(ds,
batch_size=BATCH_SIZE,
shuffle=is_shuffle,
drop_last=True,
sampler=data_sampler,
**kwargs)
logger.info("Processes {}/{} ({:.0f}%) of train data".format(
len(dl.sampler), len(dl.dataset),
100. * len(dl.sampler) / len(dl.dataset)))
# initialize DALL-E
dalle = DALLE(vae=vae, **dalle_params)
if not using_deepspeed:
if args.fp16:
dalle = dalle.half()
dalle = dalle.cuda()
if RESUME and not using_deepspeed:
dalle.load_state_dict(weights)
# optimizer
opt = Adam(get_trainable_params(dalle), lr=LEARNING_RATE)
if LR_DECAY:
scheduler = ReduceLROnPlateau(
opt,
mode="min",
factor=0.5,
patience=10,
cooldown=10,
min_lr=1e-6,
verbose=True,
)
# if distr_backend.is_root_worker():
if args.rank == 0:
# experiment tracker
model_config = dict(depth=DEPTH, heads=HEADS, dim_head=DIM_HEAD)
# wandb_dir = '/tmp/wandb'
# if not os.path.exists(wandb_dir):
# os.makedirs(wandb_dir)
run = wandb.init(
project=args.wandb_name, # 'dalle_train_transformer' by default
resume=RESUME,
config=model_config,
# dir=wandb_dir
)
# distribute
distr_backend.check_batch_size(BATCH_SIZE)
deepspeed_config = {
'train_batch_size': BATCH_SIZE,
'gradient_clipping': GRAD_CLIP_NORM,
'fp16': {
'enabled': args.fp16,
},
}
(distr_dalle, distr_opt, distr_dl,
distr_scheduler) = distr_backend.distribute(
args=args,
model=dalle,
optimizer=opt,
model_parameters=get_trainable_params(dalle),
training_data=ds if using_deepspeed else dl,
lr_scheduler=scheduler if LR_DECAY else None,
config_params=deepspeed_config,
)
avoid_model_calls = using_deepspeed and args.fp16
if RESUME and using_deepspeed:
distr_dalle.load_checkpoint(str(cp_dir))
# training
for epoch in range(EPOCHS):
logger.debug(f"********* epoch : {epoch} **********")
if data_sampler:
data_sampler.set_epoch(epoch)
for i, (text, images) in enumerate(distr_dl):
if i % 10 == 0 and args.rank == 0:
t = time.time()
if args.fp16:
images = images.half()
text, images = map(lambda t: t.cuda(), (text, images))
loss = distr_dalle(text, images, return_loss=True)
if using_deepspeed:
distr_dalle.backward(loss)
distr_dalle.step()
# Gradients are automatically zeroed after the step
else:
loss.backward()
clip_grad_norm_(distr_dalle.parameters(), GRAD_CLIP_NORM)
distr_opt.step()
distr_opt.zero_grad()
# Collective loss, averaged
avg_loss = distr_backend.average_all(loss)
log = {}
# if i % 10 == 0 and distr_backend.is_root_worker():
if i % 10 == 0 and args.rank == 0:
print(epoch, i, f'loss - {avg_loss.item()}')
log = {
**log, 'epoch': epoch,
'iter': i,
'loss': avg_loss.item()
}
if i % SAVE_EVERY_N_STEPS == 0:
args.distr_dalle = distr_dalle
args.dalle_params = dalle_params
args.vae_params = vae_params
args.using_deepspeed = using_deepspeed
args.DEEPSPEED_CP_AUX_FILENAME = DEEPSPEED_CP_AUX_FILENAME
save_model(args,
f"{args.model_dir+'/'+DALLE_OUTPUT_FILE_NAME}")
if i % 100 == 0:
# if distr_backend.is_root_worker():
if args.rank == 0:
sample_text = text[:1]
token_list = sample_text.masked_select(
sample_text != 0).tolist()
decoded_text = tokenizer.decode(token_list)
if not avoid_model_calls:
# CUDA index errors when we don't guard this
image = dalle.generate_images(
text[:1], filter_thres=0.9) # topk sampling at 0.9
log = {
**log,
}
if not avoid_model_calls:
log['image'] = wandb.Image(image, caption=decoded_text)
if i % 10 == 9 and args.rank == 0:
sample_per_sec = BATCH_SIZE * 10 / (time.time() - t)
log["sample_per_sec"] = sample_per_sec
print(epoch, i, f'sample_per_sec - {sample_per_sec}')
# if distr_backend.is_root_worker():
if args.rank == 0:
wandb.log(log)
if LR_DECAY and not using_deepspeed:
# Scheduler is automatically progressed after the step when
# using DeepSpeed.
distr_scheduler.step(loss)
args.distr_dalle = distr_dalle
args.dalle_params = dalle_params
args.vae_params = vae_params
args.using_deepspeed = using_deepspeed
args.DEEPSPEED_CP_AUX_FILENAME = DEEPSPEED_CP_AUX_FILENAME
save_model(args, f"{args.model_dir+'/'+DALLE_OUTPUT_FILE_NAME}")
# sync_local_checkpoints_to_s3(local_path=f'{args.model_dir}', s3_path='s3://lgaivision-coco-usva/Dalle_Model/tmd/')
sync_local_checkpoints_to_s3(
args.model_dir,
os.path.join(args.output_s3, args.job_name + "/temp"))
# if distr_backend.is_root_worker():
if args.rank == 0:
# save trained model to wandb as an artifact every epoch's end
model_artifact = wandb.Artifact('trained-dalle',
type='model',
metadata=dict(model_config))
import glob
print(f"************** file : {glob.glob(args.model_dir+'/*')}")
try:
print(f"wandb.run.dir : {wandb.run.dir}")
print(
f"************** file wandb: {glob.glob(wandb.run.dir+'/*')}"
)
except:
pass
model_artifact.add_file(
f"{args.model_dir+'/'+DALLE_OUTPUT_FILE_NAME}")
run.log_artifact(model_artifact)
args.distr_dalle = distr_dalle
args.dalle_params = dalle_params
args.vae_params = vae_params
args.using_deepspeed = using_deepspeed
args.DEEPSPEED_CP_AUX_FILENAME = DEEPSPEED_CP_AUX_FILENAME
resource_check(args)
save_model(args, f"{args.model_dir +'/'+DALLE_OUTPUT_FILE_NAME}")
if args.rank == 0:
# from distutils.dir_util import copy_tree
# copy_tree(f'{args.model_dir}', f'{args.model_dir_last}')
sync_local_checkpoints_to_s3(
args.model_dir,
os.path.join(args.output_s3, args.job_name + "/temp"))
resource_check(args)
# if distr_backend.is_root_worker():
if args.rank == 0:
wandb.save(f"{args.model_dir+'/'+DALLE_OUTPUT_FILE_NAME}",
base_path=args.model_dir)
model_artifact = wandb.Artifact('trained-dalle',
type='model',
metadata=dict(model_config))
model_artifact.add_file(f"{args.model_dir+'/'+DALLE_OUTPUT_FILE_NAME}")
run.log_artifact(model_artifact)
wandb.finish()
# if args.local_rank == 0 or args.local_rank == 1:
# gpu_mon_thread.kill()
distributed_utils.backend.local_barrier()
print("************************ Finished *************************")
data_sampler = torch.utils.data.distributed.DistributedSampler(
ds,
num_replicas=distr_backend.get_world_size(),
rank=distr_backend.get_rank())
else:
data_sampler = None
dl = DataLoader(ds,
batch_size=BATCH_SIZE,
shuffle=not data_sampler,
drop_last=True,
sampler=data_sampler)
# initialize DALL-E
dalle = DALLE(vae=vae, **dalle_params)
if args.fp16:
dalle = dalle.half()
dalle = dalle.cuda()
if RESUME:
dalle.load_state_dict(weights)
# optimizer
opt = Adam(dalle.parameters(), lr=LEARNING_RATE)
if LR_DECAY:
scheduler = ReduceLROnPlateau(
opt,
mode="min",
compose = T.Compose([T.Resize(IMAGE_SIZE),
T.CenterCrop(IMAGE_SIZE),
T.ToTensor(),])
def collate_fn(batch):
return tuple(zip(*batch))
ds = CocoCaptions(root=IMAGE_PATH, annFile=ANNO_PATH, transform=compose)
dl = DataLoader(ds, BATCH_SIZE, shuffle=True, num_workers=0, collate_fn=collate_fn)
assert len(ds) > 0, 'dataset is empty'
print(f'{len(ds)} image-text pairs found for training')
# initialize DALL-E
dalle = DALLE(**dalle_params)
if RESUME:
dalle.load_state_dict(weights)
dalle = torch.nn.DataParallel(dalle).cuda()
# optimizer
opt = Adam(dalle.parameters(), lr = LEARNING_RATE)
# experiment tracker
import wandb
wandb.config.depth = DEPTH
wandb.config.heads = HEADS
wandb.config.dim_head = DIM_HEAD
assert dalle_path.exists(), 'trained DALL-E must exist'
load_obj = torch.load(str(dalle_path))
dalle_params, vae_params, weights = load_obj.pop('hparams'), load_obj.pop(
'vae_params'), load_obj.pop('weights')
dalle_params.pop('vae', None) # cleanup later
if args.taming:
vae = VQGanVAE(args.vqgan_model_path, args.vqgan_config_path)
elif vae_params is not None:
vae = DiscreteVAE(**vae_params)
else:
vae = OpenAIDiscreteVAE()
dalle = DALLE(vae=vae, **dalle_params).cuda()
dalle.load_state_dict(weights)
# generate images
image_size = vae.image_size
texts = args.text.split('|')
for j, text in tqdm(enumerate(texts)):
if args.gentxt:
text_tokens, gen_texts = dalle.generate_texts(tokenizer,
text=text,
filter_thres=args.top_k)
text = gen_texts[0]