(self, path)
| 99 | self.state.iteration = iteration |
| 100 | |
| 101 | def save_checkpoint(self, path): |
| 102 | logger.info("Saving checkpoint to file {}".format(path)) |
| 103 | t_start = time.time() |
| 104 | |
| 105 | state_dict = {} |
| 106 | |
| 107 | from collections import OrderedDict |
| 108 | |
| 109 | new_state_dict = OrderedDict() |
| 110 | for k, v in self.state.model.state_dict().items(): |
| 111 | key = k |
| 112 | if k.split(".")[0] == "module": |
| 113 | key = k[7:] |
| 114 | new_state_dict[key] = v |
| 115 | state_dict["model"] = new_state_dict |
| 116 | state_dict["optimizer"] = self.state.optimizer.state_dict() |
| 117 | state_dict["epoch"] = self.state.epoch |
| 118 | state_dict["iteration"] = self.state.iteration |
| 119 | |
| 120 | t_iobegin = time.time() |
| 121 | torch.save(state_dict, path) |
| 122 | del state_dict |
| 123 | del new_state_dict |
| 124 | t_end = time.time() |
| 125 | logger.info( |
| 126 | "Save checkpoint to file {}, Time usage:\n\tprepare checkpoint: {}, IO: {}".format( |
| 127 | path, t_iobegin - t_start, t_end - t_iobegin |
| 128 | ) |
| 129 | ) |
| 130 | |
| 131 | def link_tb(self, source, target): |
| 132 | ensure_dir(source) |
no test coverage detected