def load_encoder(checkpoint_path):
"""Load the encoder portion from a model path"""
checkpoint = torch.load(checkpoint_path)
opts = checkpoint['opts']
opts['checkpoint_path'] = checkpoint_path
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts['device'] = "cuda"
opts = Namespace(**opts)
if opts.encoder_type == 'GradualStyleEncoder':
encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', opts)
elif opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoW':
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(
50, 'ir_se', opts)
elif opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoWPlus':
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoWPlus(
50, 'ir_se', opts)
else:
raise Exception('{} is not a valid encoder'.format(opts.encoder_type))
encoder.load_state_dict(get_keys(checkpoint, "encoder"), strict=True)
return encoder
def set_encoder(self):
if self.opts.encoder_type == 'GradualStyleEncoder':
encoder = psp_encoders.GradualStyleEncoder(50, 'ir_se', self.opts)
elif self.opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoW':
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(50, 'ir_se', self.opts)
elif self.opts.encoder_type == 'BackboneEncoderUsingLastLayerIntoWPlus':
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoWPlus(50, 'ir_se', self.opts)
else:
raise Exception('{} is not a valid encoders'.format(self.opts.encoder_type))
return encoder
def set_encoder(self):
if self.opts.encoder_type == "GradualStyleEncoder":
encoder = psp_encoders.GradualStyleEncoder(50, "ir_se", self.opts)
elif self.opts.encoder_type == "BackboneEncoderUsingLastLayerIntoW":
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoW(
50, "ir_se", self.opts)
elif (self.opts.encoder_type ==
"BackboneEncoderUsingLastLayerIntoWPlus"):
encoder = psp_encoders.BackboneEncoderUsingLastLayerIntoWPlus(
50, "ir_se", self.opts)
else:
raise Exception("{} is not a valid encoders".format(
self.opts.encoder_type))
return encoder