kxngluca/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Revert recent changes to extr

Browse Source
This commit is contained in:
James Betker 2021-10-30 20:48:06 -06:00
parent 83cccef9d8
commit b404a3b747
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
codes/trainer/ExtensibleTrainer.py
View File

@ -21,7 +21,7 @@ logger = logging.getLogger('base')
class ExtensibleTrainer(BaseModel): class ExtensibleTrainer(BaseModel):
def __init__(self, opt): def __init__(self, opt, cached_networks={}):
super(ExtensibleTrainer, self).__init__(opt) super(ExtensibleTrainer, self).__init__(opt)
if opt['dist']: if opt['dist']:
self.rank = torch.distributed.get_rank() self.rank = torch.distributed.get_rank()
@ -54,6 +54,10 @@ class ExtensibleTrainer(BaseModel):
if 'trainable' not in net.keys(): if 'trainable' not in net.keys():
net['trainable'] = True net['trainable'] = True
if name in cached_networks.keys():
new_net = cached_networks[name]
else:
new_net = None
if net['type'] == 'generator': if net['type'] == 'generator':
if new_net is None: if new_net is None:
new_net = networks.create_model(opt, net, self.netsG).to(self.device) new_net = networks.create_model(opt, net, self.netsG).to(self.device)