mirror of
https://github.com/lucidrains/DALLE2-pytorch.git
synced 2025-12-19 17:54:20 +01:00
remove indirection
This commit is contained in:
Phil Wang
@@ -111,37 +111,110 @@ def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,N
|
|||||||
"CosineSimilarity(text_embed,predicted_unrelated_embed)": np.mean(unrelated_similarity),
|
"CosineSimilarity(text_embed,predicted_unrelated_embed)": np.mean(unrelated_similarity),
|
||||||
"Cosine similarity difference":np.mean(predicted_similarity - original_similarity)})
|
"Cosine similarity difference":np.mean(predicted_similarity - original_similarity)})
|
||||||
|
|
||||||
def train(image_embed_dim,
|
@click.command()
|
||||||
image_embed_url,
|
@click.option("--wandb-entity", default="laion")
|
||||||
text_embed_url,
|
@click.option("--wandb-project", default="diffusion-prior")
|
||||||
batch_size,
|
@click.option("--wandb-dataset", default="LAION-5B")
|
||||||
train_percent,
|
@click.option("--wandb-arch", default="DiffusionPrior")
|
||||||
val_percent,
|
@click.option("--image-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
|
||||||
test_percent,
|
@click.option("--text-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
|
||||||
num_epochs,
|
@click.option("--learning-rate", default=1.1e-4)
|
||||||
dp_loss_type,
|
@click.option("--weight-decay", default=6.02e-2)
|
||||||
clip,
|
@click.option("--dropout", default=5e-2)
|
||||||
dp_condition_on_text_encodings,
|
@click.option("--max-grad-norm", default=0.5)
|
||||||
dp_timesteps,
|
@click.option("--batch-size", default=10**4)
|
||||||
dp_normformer,
|
@click.option("--num-epochs", default=5)
|
||||||
dp_cond_drop_prob,
|
@click.option("--image-embed-dim", default=768)
|
||||||
dpn_depth,
|
@click.option("--train-percent", default=0.7)
|
||||||
dpn_dim_head,
|
@click.option("--val-percent", default=0.2)
|
||||||
dpn_heads,
|
@click.option("--test-percent", default=0.1)
|
||||||
save_interval,
|
@click.option("--dpn-depth", default=6)
|
||||||
save_path,
|
@click.option("--dpn-dim-head", default=64)
|
||||||
device,
|
@click.option("--dpn-heads", default=8)
|
||||||
RESUME,
|
@click.option("--dp-condition-on-text-encodings", default=False)
|
||||||
DPRIOR_PATH,
|
@click.option("--dp-timesteps", default=100)
|
||||||
config,
|
@click.option("--dp-normformer", default=False)
|
||||||
wandb_entity,
|
@click.option("--dp-cond-drop-prob", default=0.1)
|
||||||
wandb_project,
|
@click.option("--dp-loss-type", default="l2")
|
||||||
learning_rate=0.001,
|
@click.option("--clip", default=None)
|
||||||
max_grad_norm=0.5,
|
@click.option("--amp", default=False)
|
||||||
weight_decay=0.01,
|
@click.option("--save-interval", default=30)
|
||||||
dropout=0.05,
|
@click.option("--save-path", default="./diffusion_prior_checkpoints")
|
||||||
amp=False):
|
@click.option("--pretrained-model-path", default=None)
|
||||||
|
def train(
|
||||||
|
wandb_entity,
|
||||||
|
wandb_project,
|
||||||
|
wandb_dataset,
|
||||||
|
wandb_arch,
|
||||||
|
image_embed_url,
|
||||||
|
text_embed_url,
|
||||||
|
learning_rate,
|
||||||
|
weight_decay,
|
||||||
|
dropout,
|
||||||
|
max_grad_norm,
|
||||||
|
batch_size,
|
||||||
|
num_epochs,
|
||||||
|
image_embed_dim,
|
||||||
|
train_percent,
|
||||||
|
val_percent,
|
||||||
|
test_percent,
|
||||||
|
dpn_depth,
|
||||||
|
dpn_dim_head,
|
||||||
|
dpn_heads,
|
||||||
|
dp_condition_on_text_encodings,
|
||||||
|
dp_timesteps,
|
||||||
|
dp_normformer,
|
||||||
|
dp_cond_drop_prob,
|
||||||
|
dp_loss_type,
|
||||||
|
clip,
|
||||||
|
amp,
|
||||||
|
save_interval,
|
||||||
|
save_path,
|
||||||
|
pretrained_model_path
|
||||||
|
):
|
||||||
|
config = {
|
||||||
|
"learning_rate": learning_rate,
|
||||||
|
"architecture": wandb_arch,
|
||||||
|
"dataset": wandb_dataset,
|
||||||
|
"weight_decay": weight_decay,
|
||||||
|
"max_gradient_clipping_norm": max_grad_norm,
|
||||||
|
"batch_size": batch_size,
|
||||||
|
"epochs": num_epochs,
|
||||||
|
"diffusion_prior_network": {
|
||||||
|
"depth": dpn_depth,
|
||||||
|
"dim_head": dpn_dim_head,
|
||||||
|
"heads": dpn_heads,
|
||||||
|
"normformer": dp_normformer
|
||||||
|
},
|
||||||
|
"diffusion_prior": {
|
||||||
|
"condition_on_text_encodings": dp_condition_on_text_encodings,
|
||||||
|
"timesteps": dp_timesteps,
|
||||||
|
"cond_drop_prob": dp_cond_drop_prob,
|
||||||
|
"loss_type": dp_loss_type,
|
||||||
|
"clip": clip
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
# Check if DPRIOR_PATH exists(saved model path)
|
||||||
|
|
||||||
|
DPRIOR_PATH = args.pretrained_model_path
|
||||||
|
RESUME = exists(DPRIOR_PATH)
|
||||||
|
|
||||||
|
if not RESUME:
|
||||||
|
tracker.init(
|
||||||
|
entity = wandb_entity,
|
||||||
|
project = wandb_project,
|
||||||
|
config = config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Obtain the utilized device.
|
||||||
|
|
||||||
|
has_cuda = torch.cuda.is_available()
|
||||||
|
if has_cuda:
|
||||||
|
device = torch.device("cuda:0")
|
||||||
|
torch.cuda.set_device(device)
|
||||||
|
|
||||||
|
# Training loop
|
||||||
# diffusion prior network
|
# diffusion prior network
|
||||||
|
|
||||||
prior_network = DiffusionPriorNetwork(
|
prior_network = DiffusionPriorNetwork(
|
||||||
@@ -269,140 +342,5 @@ def train(image_embed_dim,
|
|||||||
end = num_data_points
|
end = num_data_points
|
||||||
eval_model(diffusion_prior,device,image_reader,text_reader,start,end,batch_size,dp_loss_type,phase="Test")
|
eval_model(diffusion_prior,device,image_reader,text_reader,start,end,batch_size,dp_loss_type,phase="Test")
|
||||||
|
|
||||||
@click.command()
|
|
||||||
@click.option("--wandb-entity", default="laion")
|
|
||||||
@click.option("--wandb-project", default="diffusion-prior")
|
|
||||||
@click.option("--wandb-dataset", default="LAION-5B")
|
|
||||||
@click.option("--wandb-arch", default="DiffusionPrior")
|
|
||||||
@click.option("--image-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/img_emb/")
|
|
||||||
@click.option("--text-embed-url", default="https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion2B-en/text_emb/")
|
|
||||||
@click.option("--learning-rate", default=1.1e-4)
|
|
||||||
@click.option("--weight-decay", default=6.02e-2)
|
|
||||||
@click.option("--dropout", default=5e-2)
|
|
||||||
@click.option("--max-grad-norm", default=0.5)
|
|
||||||
@click.option("--batch-size", default=10**4)
|
|
||||||
@click.option("--num-epochs", default=5)
|
|
||||||
@click.option("--image-embed-dim", default=768)
|
|
||||||
@click.option("--train-percent", default=0.7)
|
|
||||||
@click.option("--val-percent", default=0.2)
|
|
||||||
@click.option("--test-percent", default=0.1)
|
|
||||||
@click.option("--dpn-depth", default=6)
|
|
||||||
@click.option("--dpn-dim-head", default=64)
|
|
||||||
@click.option("--dpn-heads", default=8)
|
|
||||||
@click.option("--dp-condition-on-text-encodings", default=False)
|
|
||||||
@click.option("--dp-timesteps", default=100)
|
|
||||||
@click.option("--dp-normformer", default=False)
|
|
||||||
@click.option("--dp-cond-drop-prob", default=0.1)
|
|
||||||
@click.option("--dp-loss-type", default="l2")
|
|
||||||
@click.option("--clip", default=None)
|
|
||||||
@click.option("--amp", default=False)
|
|
||||||
@click.option("--save-interval", default=30)
|
|
||||||
@click.option("--save-path", default="./diffusion_prior_checkpoints")
|
|
||||||
@click.option("--pretrained-model-path", default=None)
|
|
||||||
def main(
|
|
||||||
wandb_entity,
|
|
||||||
wandb_project,
|
|
||||||
wandb_dataset,
|
|
||||||
wandb_arch,
|
|
||||||
image_embed_url,
|
|
||||||
text_embed_url,
|
|
||||||
learning_rate,
|
|
||||||
weight_decay,
|
|
||||||
dropout,
|
|
||||||
max_grad_norm,
|
|
||||||
batch_size,
|
|
||||||
num_epochs,
|
|
||||||
image_embed_dim,
|
|
||||||
train_percent,
|
|
||||||
val_percent,
|
|
||||||
test_percent,
|
|
||||||
dpn_depth,
|
|
||||||
dpn_dim_head,
|
|
||||||
dpn_heads,
|
|
||||||
dp_condition_on_text_encodings,
|
|
||||||
dp_timesteps,
|
|
||||||
dp_normformer,
|
|
||||||
dp_cond_drop_prob,
|
|
||||||
dp_loss_type,
|
|
||||||
clip,
|
|
||||||
amp,
|
|
||||||
save_interval,
|
|
||||||
save_path,
|
|
||||||
pretrained_model_path
|
|
||||||
):
|
|
||||||
config = {
|
|
||||||
"learning_rate": learning_rate,
|
|
||||||
"architecture": wandb_arch,
|
|
||||||
"dataset": wandb_dataset,
|
|
||||||
"weight_decay": weight_decay,
|
|
||||||
"max_gradient_clipping_norm": max_grad_norm,
|
|
||||||
"batch_size": batch_size,
|
|
||||||
"epochs": num_epochs,
|
|
||||||
"diffusion_prior_network": {
|
|
||||||
"depth": dpn_depth,
|
|
||||||
"dim_head": dpn_dim_head,
|
|
||||||
"heads": dpn_heads,
|
|
||||||
"normformer": dp_normformer
|
|
||||||
},
|
|
||||||
"diffusion_prior": {
|
|
||||||
"condition_on_text_encodings": dp_condition_on_text_encodings,
|
|
||||||
"timesteps": dp_timesteps,
|
|
||||||
"cond_drop_prob": dp_cond_drop_prob,
|
|
||||||
"loss_type": dp_loss_type,
|
|
||||||
"clip": clip
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
# Check if DPRIOR_PATH exists(saved model path)
|
|
||||||
|
|
||||||
DPRIOR_PATH = args.pretrained_model_path
|
|
||||||
RESUME = exists(DPRIOR_PATH)
|
|
||||||
|
|
||||||
if not RESUME:
|
|
||||||
tracker.init(
|
|
||||||
entity = wandb_entity,
|
|
||||||
project = wandb_project,
|
|
||||||
config = config
|
|
||||||
)
|
|
||||||
|
|
||||||
# Obtain the utilized device.
|
|
||||||
|
|
||||||
has_cuda = torch.cuda.is_available()
|
|
||||||
if has_cuda:
|
|
||||||
device = torch.device("cuda:0")
|
|
||||||
torch.cuda.set_device(device)
|
|
||||||
|
|
||||||
# Training loop
|
|
||||||
train(image_embed_dim,
|
|
||||||
image_embed_url,
|
|
||||||
text_embed_url,
|
|
||||||
batch_size,
|
|
||||||
train_percent,
|
|
||||||
val_percent,
|
|
||||||
test_percent,
|
|
||||||
num_epochs,
|
|
||||||
dp_loss_type,
|
|
||||||
clip,
|
|
||||||
dp_condition_on_text_encodings,
|
|
||||||
dp_timesteps,
|
|
||||||
dp_normformer,
|
|
||||||
dp_cond_drop_prob,
|
|
||||||
dpn_depth,
|
|
||||||
dpn_dim_head,
|
|
||||||
dpn_heads,
|
|
||||||
save_interval,
|
|
||||||
save_path,
|
|
||||||
device,
|
|
||||||
RESUME,
|
|
||||||
DPRIOR_PATH,
|
|
||||||
config,
|
|
||||||
wandb_entity,
|
|
||||||
wandb_project,
|
|
||||||
learning_rate,
|
|
||||||
max_grad_norm,
|
|
||||||
weight_decay,
|
|
||||||
dropout,
|
|
||||||
amp)
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
train()
|
||||||
|
|||||||
Reference in New Issue
Block a user