(self, step, is_best)
| 938 | # save the trained generator, generator_ema, and discriminator. |
| 939 | # ----------------------------------------------------------------------------- |
| 940 | def save(self, step, is_best): |
| 941 | when = "best" if is_best is True else "current" |
| 942 | misc.make_GAN_untrainable(self.Gen, self.Gen_ema, self.Dis) |
| 943 | Gen, Gen_ema, Dis = misc.peel_models(self.Gen, self.Gen_ema, self.Dis) |
| 944 | |
| 945 | g_states = {"state_dict": Gen.state_dict(), "optimizer": self.OPTIMIZATION.g_optimizer.state_dict()} |
| 946 | |
| 947 | d_states = { |
| 948 | "state_dict": Dis.state_dict(), |
| 949 | "optimizer": self.OPTIMIZATION.d_optimizer.state_dict(), |
| 950 | "seed": self.RUN.seed, |
| 951 | "run_name": self.run_name, |
| 952 | "step": step, |
| 953 | "epoch": self.epoch_counter, |
| 954 | "topk": self.topk, |
| 955 | "aa_p": self.aa_p, |
| 956 | "best_step": self.best_step, |
| 957 | "best_fid": self.best_fid, |
| 958 | "best_fid_ckpt": self.RUN.ckpt_dir, |
| 959 | "lecam_emas": self.lecam_ema.__dict__, |
| 960 | } |
| 961 | |
| 962 | if self.Gen_ema is not None: |
| 963 | g_ema_states = {"state_dict": Gen_ema.state_dict()} |
| 964 | |
| 965 | misc.save_model(model="G", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_states) |
| 966 | misc.save_model(model="D", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=d_states) |
| 967 | if self.Gen_ema is not None: |
| 968 | misc.save_model(model="G_ema", when=when, step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_ema_states) |
| 969 | |
| 970 | if when == "best": |
| 971 | misc.save_model(model="G", when="current", step=step, ckpt_dir=self.RUN.ckpt_dir, states=g_states) |
| 972 | misc.save_model(model="D", when="current", step=step, ckpt_dir=self.RUN.ckpt_dir, states=d_states) |
| 973 | if self.Gen_ema is not None: |
| 974 | misc.save_model(model="G_ema", |
| 975 | when="current", |
| 976 | step=step, |
| 977 | ckpt_dir=self.RUN.ckpt_dir, |
| 978 | states=g_ema_states) |
| 979 | |
| 980 | if self.global_rank == 0 and self.logger: |
| 981 | self.logger.info("Save model to {}".format(self.RUN.ckpt_dir)) |
| 982 | |
| 983 | misc.make_GAN_trainable(self.Gen, self.Gen_ema, self.Dis) |
| 984 | |
| 985 | # ----------------------------------------------------------------------------- |
| 986 | # save real images to measure metrics for evaluation. |
no outgoing calls
no test coverage detected