def __init__(self, cfg, outputs_desc):
super().__init__()
print('outputs_desc {}'.format(outputs_desc))
self.outputs_desc = outputs_desc
ch_out = sum(outputs_desc.values())
self.encoder = Encoder(
cfg.model_encoder_name,
pretrained=True, ########CHANGE TO TRUE WHEN UPLOADING
zero_init_residual=True,
replace_stride_with_dilation=(False, False, True),
)
print('USING DISTILLATION MODEL :DDDDDDDDDDDDDDD')
ch_out_encoder_bottleneck, ch_out_encoder_4x = get_encoder_channel_counts(
cfg.model_encoder_name)
# ch_out_encoder_4x = 512
print('ch out encoder 4x: ', ch_out_encoder_4x)
self.aspp_semseg = ASPP(ch_out_encoder_bottleneck, 256)
self.decoder_semseg = DecoderDeeplabV3p(256, ch_out_encoder_4x,
ch_out - 1)
self.assp_depth = ASPP(ch_out_encoder_bottleneck, 256)
self.decoder_depth = DecoderDeeplabV3p(256, ch_out_encoder_4x, 1)
self.sa_from_depth2semseg = SelfAttention(256, 256)
self.sa_from_semseg2depth = SelfAttention(256, 256)
self.decoder_semseg2 = DecoderDistillation(256, ch_out - 1)
self.decoder_depth2 = DecoderDistillation(256, 1)
def __init__(self, cfg, outputs_desc):
super().__init__()
self.outputs_desc = outputs_desc
ch_out = sum(outputs_desc.values())
self.encoder = Encoder(
cfg.model_encoder_name,
pretrained=True,
zero_init_residual=True,
replace_stride_with_dilation=(False, False, True),
)
print('Branched Model')
ch_out_encoder_bottleneck, ch_out_encoder_4x = get_encoder_channel_counts(
cfg.model_encoder_name)
self.aspp_semseg = ASPP(ch_out_encoder_bottleneck, 256)
self.decoder_semseg = DecoderDeeplabV3p(256, ch_out_encoder_4x,
ch_out - 1)
self.aspp_depth = ASPP(ch_out_encoder_bottleneck, 256)
self.decoder_depth = DecoderDeeplabV3p(256, ch_out_encoder_4x, 1)