CalvinBackup/neuralchen-SimSwap
1
0
Fork 0
Code Issues 288 Pull Requests 16 Actions Packages Projects Releases 2 Wiki Activity

training scripts released

Browse Source
This commit is contained in:
chenxuanhong
2022-04-20 18:36:26 +08:00
parent 9492873690
commit f48dc8cf62
16 changed files with 1688 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/base_model.py
+54
View File
@@ -37,6 +37,19 @@ class BaseModel(torch.nn.Module):
def save(self, label):
pass
# helper saving function that can be used by subclasses
def save_network(self, network, network_label, epoch_label, gpu_ids=None):
save_filename = '{}_net_{}.pth'.format(epoch_label, network_label)
save_path = os.path.join(self.save_dir, save_filename)
torch.save(network.cpu().state_dict(), save_path)
if torch.cuda.is_available():
network.cuda()
def save_optim(self, network, network_label, epoch_label, gpu_ids=None):
save_filename = '{}_optim_{}.pth'.format(epoch_label, network_label)
save_path = os.path.join(self.save_dir, save_filename)
torch.save(network.state_dict(), save_path)
# helper saving function that can be used by subclasses
def save_network(self, network, network_label, epoch_label, gpu_ids):
@@ -63,6 +76,47 @@ class BaseModel(torch.nn.Module):
except:
pretrained_dict = torch.load(save_path)
model_dict = network.state_dict()
try:
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
network.load_state_dict(pretrained_dict)
if self.opt.verbose:
print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
except:
print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
for k, v in pretrained_dict.items():
if v.size() == model_dict[k].size():
model_dict[k] = v
if sys.version_info >= (3,0):
not_initialized = set()
else:
from sets import Set
not_initialized = Set()
for k, v in model_dict.items():
if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
not_initialized.add(k.split('.')[0])
print(sorted(not_initialized))
network.load_state_dict(model_dict)
# helper loading function that can be used by subclasses
def load_optim(self, network, network_label, epoch_label, save_dir=''):
save_filename = '%s_optim_%s.pth' % (epoch_label, network_label)
if not save_dir:
save_dir = self.save_dir
save_path = os.path.join(save_dir, save_filename)
if not os.path.isfile(save_path):
print('%s not exists yet!' % save_path)
if network_label == 'G':
raise('Generator must exist!')
else:
#network.load_state_dict(torch.load(save_path))
try:
network.load_state_dict(torch.load(save_path, map_location=torch.device("cpu")))
except:
pretrained_dict = torch.load(save_path, map_location=torch.device("cpu"))
model_dict = network.state_dict()
try:
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
network.load_state_dict(pretrained_dict)