Spaces:
Runtime error
Runtime error
File size: 3,813 Bytes
6e601ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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)
|