Merge pull request #8 from IDEALLab/feature/seed_artifacts

Tag artifacts with seed

Author: ffelten

engiopt/cgan_1d.py

@@ -291,12 +291,12 @@ def sample_designs(n_designs: int) -> th.Tensor:
 
                         th.save(ckpt_gen, "generator.pth")
                         th.save(ckpt_disc, "discriminator.pth")
-                        artifact_gen = wandb.Artifact("generator", type="model")
+                        artifact_gen = wandb.Artifact(f"{args.algo}_generator", type="model")
                         artifact_gen.add_file("generator.pth")
-                        artifact_disc = wandb.Artifact("discriminator", type="model")
+                        artifact_disc = wandb.Artifact(f"{args.algo}_discriminator", type="model")
                         artifact_disc.add_file("discriminator.pth")
 
-                        wandb.log_artifact(artifact_gen)
-                        wandb.log_artifact(artifact_disc)
+                        wandb.log_artifact(artifact_gen, aliases=[f"seed_{args.seed}"])
+                        wandb.log_artifact(artifact_disc, aliases=[f"seed_{args.seed}"])
 
     wandb.finish()

engiopt/diffusion_1d.py

@@ -175,8 +175,9 @@ class Args:
                         }
 
                         th.save(ckpt, "model.pth")
-                        artifact = wandb.Artifact("diffusion", type="model")
+                        artifact = wandb.Artifact(f"{args.algo}_model", type="model")
                         artifact.add_file("model.pth")
-                        wandb.log_artifact(artifact)
+
+                        wandb.log_artifact(artifact, aliases=[f"seed_{args.seed}"])
 
     wandb.finish()

engiopt/gan_1d.py

@@ -261,12 +261,12 @@ def forward(self, design: th.Tensor) -> th.Tensor:  # noqa: D102
 
                         th.save(ckpt_gen, "generator.pth")
                         th.save(ckpt_disc, "discriminator.pth")
-                        artifact_gen = wandb.Artifact("generator", type="model")
+                        artifact_gen = wandb.Artifact(f"{args.algo}_generator", type="model")
                         artifact_gen.add_file("generator.pth")
-                        artifact_disc = wandb.Artifact("discriminator", type="model")
+                        artifact_disc = wandb.Artifact(f"{args.algo}_discriminator", type="model")
                         artifact_disc.add_file("discriminator.pth")
 
-                        wandb.log_artifact(artifact_gen)
-                        wandb.log_artifact(artifact_disc)
+                        wandb.log_artifact(artifact_gen, aliases=[f"seed_{args.seed}"])
+                        wandb.log_artifact(artifact_disc, aliases=[f"seed_{args.seed}"])
 
     wandb.finish()