def __init__(self, hparams):
super(RoadMapNetwork, self).__init__()
self.hparams = hparams
dropout = False if self.hparams.DROPOUT == 0 else self.hparams.DROPOUT
self.apply_sigmoid = True
if self.hparams.LOSS in ["bce", "weighted_bce"]:
self.apply_sigmoid = False
self.loss_fn = LOSS[self.hparams.LOSS]
self.feature_extractor = DenoisingAutoencoder.load_from_checkpoint(
FEATURE_EXTRACTOR_PATH) # Output size -> (None, 192, 13, 13)
self.feature_extractor.freeze()
self.classifier = UNet(
num_layers=self.hparams.NUM_LAYERS,
features_start=self.hparams.FEATURES_START,
dropout=dropout,
)
def load_pretrained_layers(self):
# Current state of base
state_dict = self.state_dict()
param_names = list(state_dict.keys())
# Load the pre-trained autoencoder encoder layer
model = DenoisingAutoencoder.load_from_checkpoint('denoising.ckpt')
temp_enc = list(model.children())[:-1]
pretrained_state_dict = temp_enc[0].state_dict()
pretrained_param_names = list(pretrained_state_dict.keys())
#We update the first 14 parameters of the base model.
for i, param in enumerate(param_names[:14]):
state_dict[param] = pretrained_state_dict[
pretrained_param_names[i]]
self.load_state_dict(state_dict)
print("\nLoaded base model.\n")
