load saved checkpoint
(self, name)
| 157 | self.net.to(self.device) |
| 158 | |
| 159 | def load_ckpt(self, name): |
| 160 | """load saved checkpoint""" |
| 161 | load_path = os.path.join(self.log_dir, f"ckpt_{name}.pth") |
| 162 | checkpoint = torch.load(load_path, map_location=self.device) |
| 163 | |
| 164 | self.net.load_state_dict(checkpoint["net"]) |
| 165 | self.Ka = checkpoint["Ka"] |
| 166 | self.Kd = checkpoint["Kd"] |
| 167 | self.Ks = checkpoint["Ks"] |
| 168 | self.Ns = checkpoint["Ns"] |
| 169 | self.aabb = torch.tensor(checkpoint["aabb"], dtype=torch.float32, device=self.device) |
| 170 | self.featmap_size = checkpoint["featmap_size"] |
| 171 | self.net.reset_aabb(self.aabb) |
| 172 | |
| 173 | if hasattr(self, "optimizer"): |
| 174 | self.optimizer.load_state_dict(checkpoint["optimizer"]) |
| 175 | self.scheduler.load_state_dict(checkpoint["scheduler"]) |
| 176 | self.net.to(self.device) |
| 177 | |
| 178 | def update_network(self, loss_dict): |
| 179 | """update network by back propagation""" |