| 48 | |
| 49 | |
| 50 | class ModelBuilder: |
| 51 | # custom weights initialization |
| 52 | @staticmethod |
| 53 | def weights_init(m): |
| 54 | classname = m.__class__.__name__ |
| 55 | if classname.find('Conv') != -1: |
| 56 | nn.init.kaiming_normal_(m.weight.data) |
| 57 | elif classname.find('BatchNorm') != -1: |
| 58 | m.weight.data.fill_(1.) |
| 59 | m.bias.data.fill_(1e-4) |
| 60 | |
| 61 | @staticmethod |
| 62 | def build_encoder(arch='resnet50dilated', fc_dim=512, weights=''): |
| 63 | pretrained = True if len(weights) == 0 else False |
| 64 | arch = arch.lower() |
| 65 | if arch == 'resnet18dilated': |
| 66 | orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained) |
| 67 | net_encoder = ResnetDilated(orig_resnet, dilate_scale=8) |
| 68 | elif arch == 'resnet50dilated': |
| 69 | orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained) |
| 70 | net_encoder = ResnetDilated(orig_resnet, dilate_scale=8) |
| 71 | else: |
| 72 | raise Exception('Architecture undefined!') |
| 73 | |
| 74 | if len(weights) > 0: |
| 75 | print('Loading weights for net_encoder') |
| 76 | net_encoder.load_state_dict( |
| 77 | torch.load(weights, map_location=lambda storage, loc: storage), strict=False) |
| 78 | return net_encoder |
| 79 | |
| 80 | @staticmethod |
| 81 | def build_decoder(arch='ppm', |
| 82 | fc_dim=512, num_class=150, |
| 83 | weights='', use_softmax=False): |
| 84 | arch = arch.lower() |
| 85 | if arch == 'ppm': |
| 86 | net_decoder = PPM( |
| 87 | num_class=num_class, |
| 88 | fc_dim=fc_dim, |
| 89 | use_softmax=use_softmax) |
| 90 | else: |
| 91 | raise Exception('Architecture undefined!') |
| 92 | |
| 93 | net_decoder.apply(ModelBuilder.weights_init) |
| 94 | if len(weights) > 0: |
| 95 | print('Loading weights for net_decoder') |
| 96 | net_decoder.load_state_dict( |
| 97 | torch.load(weights, map_location=lambda storage, loc: storage), strict=False) |
| 98 | return net_decoder |
| 99 | |
| 100 | |
| 101 | def conv3x3_bn_relu(in_planes, out_planes, stride=1): |
nothing calls this directly
no outgoing calls
no test coverage detected