base_model.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. import os
  2. import torch
  3. import sys
  4. class BaseModel(torch.nn.Module):
  5. def name(self):
  6. return 'BaseModel'
  7. def initialize(self, opt):
  8. self.opt = opt
  9. self.gpu_ids = opt.gpu_ids
  10. self.isTrain = opt.isTrain
  11. self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
  12. self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
  13. def set_input(self, input):
  14. self.input = input
  15. def forward(self):
  16. pass
  17. # used in test time, no backprop
  18. def test(self):
  19. pass
  20. def get_image_paths(self):
  21. pass
  22. def optimize_parameters(self):
  23. pass
  24. def get_current_visuals(self):
  25. return self.input
  26. def get_current_errors(self):
  27. return {}
  28. def save(self, label):
  29. pass
  30. # helper saving function that can be used by subclasses
  31. def save_network(self, network, network_label, epoch_label, gpu_ids):
  32. save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
  33. save_path = os.path.join(self.save_dir, save_filename)
  34. torch.save(network.cpu().state_dict(), save_path)
  35. if len(gpu_ids) and torch.cuda.is_available():
  36. network.cuda()
  37. # helper loading function that can be used by subclasses
  38. def load_network(self, network, network_label, epoch_label, save_dir=''):
  39. save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
  40. if not save_dir:
  41. save_dir = self.save_dir
  42. save_path = os.path.join(save_dir, save_filename)
  43. if not os.path.isfile(save_path):
  44. print('%s not exists yet!' % save_path)
  45. if network_label == 'G':
  46. raise('Generator must exist!')
  47. else:
  48. #network.load_state_dict(torch.load(save_path))
  49. try:
  50. network.load_state_dict(torch.load(save_path))
  51. except:
  52. pretrained_dict = torch.load(save_path)
  53. model_dict = network.state_dict()
  54. try:
  55. pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
  56. network.load_state_dict(pretrained_dict)
  57. if self.opt.verbose:
  58. print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
  59. except:
  60. print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
  61. for k, v in pretrained_dict.items():
  62. if v.size() == model_dict[k].size():
  63. model_dict[k] = v
  64. if sys.version_info >= (3,0):
  65. not_initialized = set()
  66. else:
  67. from sets import Set
  68. not_initialized = Set()
  69. for k, v in model_dict.items():
  70. if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
  71. not_initialized.add(k.split('.')[0])
  72. print(sorted(not_initialized))
  73. network.load_state_dict(model_dict)
  74. def update_learning_rate():
  75. pass