def main(opt):
config = decode_config(opt.config_str)
if opt.model == 'mobile_resnet':
from models.modules.resnet_architecture.mobile_resnet_generator import MobileResnetGenerator as SuperModel
from models.modules.resnet_architecture.sub_mobile_resnet_generator import SubMobileResnetGenerator as SubModel
input_nc, output_nc = opt.input_nc, opt.output_nc
super_model = SuperModel(input_nc,
output_nc,
ngf=opt.ngf,
norm_layer=nn.InstanceNorm2d,
n_blocks=9)
sub_model = SubModel(input_nc,
output_nc,
config=config,
norm_layer=nn.InstanceNorm2d,
n_blocks=9)
elif opt.model == 'mobile_spade':
from models.modules.spade_architecture.mobile_spade_generator import MobileSPADEGenerator as SuperModel
from models.modules.spade_architecture.sub_mobile_spade_generator import SubMobileSPADEGenerator as SubModel
opt.norm_G = 'spadesyncbatch3x3'
opt.num_upsampling_layers = 'more'
opt.semantic_nc = opt.input_nc + (1 if opt.contain_dontcare_label else
0) + (0 if opt.no_instance else 1)
super_model = SuperModel(opt)
sub_model = SubModel(opt, config)
else:
raise NotImplementedError('Unknown architecture [%s]!' % opt.model)
load_network(super_model, opt.input_path)
transfer_weight(super_model, sub_model)
output_dir = os.path.dirname(opt.output_path)
os.makedirs(output_dir, exist_ok=True)
torch.save(sub_model.state_dict(), opt.output_path)
print('Successfully export the subnet at [%s].' % opt.output_path)
def main(opt):
if opt.model == 'mobile_resnet':
from models.modules.resnet_architecture.mobile_resnet_generator import MobileResnetGenerator as SuperModel
from models.modules.resnet_architecture.sub_mobile_resnet_generator import SubMobileResnetGenerator as SubModel
elif opt.model == 'mobile_spade':
# TODO
raise NotImplementedError
else:
raise NotImplementedError('Unknown architecture [%s]!' % opt.model)
config = decode_config(opt.config_str)
input_nc, output_nc = opt.input_nc, opt.output_nc
super_model = SuperModel(input_nc,
output_nc,
ngf=opt.ngf,
norm_layer=nn.InstanceNorm2d,
n_blocks=9)
sub_model = SubModel(input_nc,
output_nc,
config=config,
norm_layer=nn.InstanceNorm2d,
n_blocks=9)
load_network(super_model, opt.input_path)
transfer_weight(super_model, sub_model)
output_dir = os.path.dirname(opt.output_path)
os.makedirs(output_dir, exist_ok=True)
torch.save(sub_model.state_dict(), opt.output_path)
print('Successfully export the subnet at [%s].' % opt.output_path)
def define_G(input_nc,
output_nc,
ngf,
netG,
norm='batch',
dropout_rate=0,
init_type='normal',
init_gain=0.02,
gpu_ids=[],
opt=None):
norm_layer = get_norm_layer(norm_type=norm)
if netG == 'resnet_9blocks':
from models.modules.resnet_architecture.resnet_generator import ResnetGenerator
net = ResnetGenerator(input_nc,
output_nc,
ngf,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
n_blocks=9)
elif netG == 'mobile_resnet_9blocks':
from models.modules.resnet_architecture.mobile_resnet_generator import MobileResnetGenerator
net = MobileResnetGenerator(input_nc,
output_nc,
ngf=ngf,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
n_blocks=9)
elif netG == 'super_mobile_resnet_9blocks':
from models.modules.resnet_architecture.super_mobile_resnet_generator import SuperMobileResnetGenerator
net = SuperMobileResnetGenerator(input_nc,
output_nc,
ngf=ngf,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
n_blocks=9)
elif netG == 'sub_mobile_resnet_9blocks':
from models.modules.resnet_architecture.sub_mobile_resnet_generator import SubMobileResnetGenerator
assert opt.config_str is not None
config = decode_config(opt.config_str)
net = SubMobileResnetGenerator(input_nc,
output_nc,
config,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
n_blocks=9)
else:
raise NotImplementedError(
'Generator model name [%s] is not recognized' % netG)
return init_net(net, init_type, init_gain, gpu_ids)