| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| from climategan.deeplab.deeplab_v2 import DeepLabV2Decoder |
| from climategan.deeplab.deeplab_v3 import DeepLabV3Decoder |
| from climategan.deeplab.mobilenet_v3 import MobileNetV2 |
| from climategan.deeplab.resnet101_v3 import ResNet101 |
| from climategan.deeplab.resnetmulti_v2 import ResNetMulti |
|
|
|
|
| def create_encoder(opts, no_init=False, verbose=0): |
| if opts.gen.encoder.architecture == "deeplabv2": |
| if verbose > 0: |
| print(" - Add Deeplabv2 Encoder") |
| return DeeplabV2Encoder(opts, no_init, verbose) |
| elif opts.gen.encoder.architecture == "deeplabv3": |
| if verbose > 0: |
| backone = opts.gen.deeplabv3.backbone |
| print(" - Add Deeplabv3 ({}) Encoder".format(backone)) |
| return build_v3_backbone(opts, no_init) |
| else: |
| raise NotImplementedError( |
| "Unknown encoder: {}".format(opts.gen.encoder.architecture) |
| ) |
|
|
|
|
| def create_segmentation_decoder(opts, no_init=False, verbose=0): |
| if opts.gen.s.architecture == "deeplabv2": |
| if verbose > 0: |
| print(" - Add DeepLabV2Decoder") |
| return DeepLabV2Decoder(opts) |
| elif opts.gen.s.architecture == "deeplabv3": |
| if verbose > 0: |
| print(" - Add DeepLabV3Decoder") |
| return DeepLabV3Decoder(opts, no_init) |
| else: |
| raise NotImplementedError( |
| "Unknown Segmentation architecture: {}".format(opts.gen.s.architecture) |
| ) |
|
|
|
|
| def build_v3_backbone(opts, no_init, verbose=0): |
| backbone = opts.gen.deeplabv3.backbone |
| output_stride = opts.gen.deeplabv3.output_stride |
| if backbone == "resnet": |
| resnet = ResNet101( |
| output_stride=output_stride, |
| BatchNorm=nn.BatchNorm2d, |
| verbose=verbose, |
| no_init=no_init, |
| ) |
| if not no_init: |
| if opts.gen.deeplabv3.backbone == "resnet": |
| assert Path(opts.gen.deeplabv3.pretrained_model.resnet).exists() |
|
|
| std = torch.load(opts.gen.deeplabv3.pretrained_model.resnet) |
| resnet.load_state_dict( |
| { |
| k.replace("backbone.", ""): v |
| for k, v in std.items() |
| if k.startswith("backbone.") |
| } |
| ) |
| print( |
| " - Loaded pre-trained DeepLabv3+ Resnet101 Backbone as Encoder" |
| ) |
| return resnet |
|
|
| elif opts.gen.deeplabv3.backbone == "mobilenet": |
| assert Path(opts.gen.deeplabv3.pretrained_model.mobilenet).exists() |
| mobilenet = MobileNetV2( |
| no_init=no_init, |
| pretrained_path=opts.gen.deeplabv3.pretrained_model.mobilenet, |
| ) |
| print(" - Loaded pre-trained DeepLabv3+ MobileNetV2 Backbone as Encoder") |
| return mobilenet |
|
|
| else: |
| raise NotImplementedError("Unknown backbone in " + str(opts.gen.deeplabv3)) |
|
|
|
|
| class DeeplabV2Encoder(nn.Module): |
| def __init__(self, opts, no_init=False, verbose=0): |
| """Deeplab architecture encoder""" |
| super().__init__() |
|
|
| self.model = ResNetMulti(opts.gen.deeplabv2.nblocks, opts.gen.encoder.n_res) |
| if opts.gen.deeplabv2.use_pretrained and not no_init: |
| saved_state_dict = torch.load(opts.gen.deeplabv2.pretrained_model) |
| new_params = self.model.state_dict().copy() |
| for i in saved_state_dict: |
| i_parts = i.split(".") |
| if not i_parts[1] in ["layer5", "resblock"]: |
| new_params[".".join(i_parts[1:])] = saved_state_dict[i] |
| self.model.load_state_dict(new_params) |
| if verbose > 0: |
| print(" - Loaded pretrained weights") |
|
|
| def forward(self, x): |
| return self.model(x) |
|
|
