Load the model saved at the filenames. It automatically gathers all the nn.modules and optimizers and load them. Everything will be loaded onto cpu.
(self, filename: str, device: torch.device = torch.device("cpu"))
| 645 | v.eval() |
| 646 | |
| 647 | def load(self, filename: str, device: torch.device = torch.device("cpu")): |
| 648 | """ |
| 649 | Load the model saved at the filenames. |
| 650 | It automatically gathers all the nn.modules and optimizers and load them. |
| 651 | Everything will be loaded onto cpu. |
| 652 | """ |
| 653 | checkpoint = torch.load(filename, map_location=device) # dict of [name, data] |
| 654 | loaded_names = set() |
| 655 | |
| 656 | # gather all base models |
| 657 | nets = inspect.getmembers(self, lambda x: isinstance(x, BaseModel)) # name, net |
| 658 | for name, net in nets: |
| 659 | if name in checkpoint: |
| 660 | getattr(self, name).load_state_dict(checkpoint[name]) |
| 661 | else: |
| 662 | warnings.warn(f"{name} not exist in checkpoint") |
| 663 | loaded_names.add(name) |
| 664 | |
| 665 | # gather all nn.modules |
| 666 | nets = inspect.getmembers(self, lambda x: isinstance(x, torch.nn.Module)) # name, net |
| 667 | for name, net in nets: |
| 668 | if name in checkpoint: |
| 669 | try: |
| 670 | getattr(self, name).load_state_dict(checkpoint[name]) |
| 671 | except: |
| 672 | new_state_dict = OrderedDict() |
| 673 | for key, val in checkpoint[name].items(): |
| 674 | if key.startswith("module."): |
| 675 | new_key = key[7:] |
| 676 | new_state_dict[new_key] = val |
| 677 | else: |
| 678 | new_state_dict[key] = val |
| 679 | getattr(self, name).load_state_dict(new_state_dict) |
| 680 | else: |
| 681 | warnings.warn(f"{name} not exist in checkpoint") |
| 682 | loaded_names.add(name) |
| 683 | |
| 684 | # gather all optimizers |
| 685 | optimizers = inspect.getmembers(self, lambda x: isinstance(x, torch.optim.Optimizer)) # name, net |
| 686 | for name, optimizer in optimizers: |
| 687 | if name in checkpoint: |
| 688 | getattr(self, name).load_state_dict(checkpoint[name]) |
| 689 | else: |
| 690 | warnings.warn(f"{name} not exist in checkpoint") |
| 691 | loaded_names.add(name) |
| 692 | |
| 693 | # gather all AMP scaler |
| 694 | try: |
| 695 | scalers = inspect.getmembers(self, lambda x: isinstance(x, torch.cuda.amp.GradScaler)) # name, net |
| 696 | for name, scaler in scalers: |
| 697 | if name in checkpoint: |
| 698 | getattr(self, name).load_state_dict(checkpoint[name]) |
| 699 | else: |
| 700 | warnings.warn(f"{name} not exist in checkpoint") |
| 701 | loaded_names.add(name) |
| 702 | except: |
| 703 | pass |
| 704 |