예제 #1
0
        keep_scale = False
        reverse_channels = False
    if args.dataset == 'voc':
        num_classes = num_classes_voc
        input_sizes = sizes_voc
        categories = categories_voc
        colors = colors_voc
    elif args.dataset == 'city':
        num_classes = num_classes_city
        input_sizes = sizes_city
        categories = categories_city
        colors = colors_city
    else:
        raise ValueError

    net = deeplab_v2(num_classes=num_classes)
    print(device)
    net.to(device)

    # Define optimizer
    # Use different learning rates if you want, we do not observe improvement from different learning rates
    params_to_optimize = [
        {
            "params":
            [p for p in net.backbone.parameters() if p.requires_grad]
        },
        {
            "params":
            [p for p in net.classifier.parameters() if p.requires_grad]
        },
    ]
예제 #2
0
import torch
from models.segmentation.segmentation import deeplab_v2
from utils.common import save_checkpoint

# COCO pre-trained Deeplab-ResNet-101 from Hung et al.
# (checked to be the same as the original caffe model published by the DeeplabV2 authors)
# (_-_) Gave them a star on github to show our gratitude
# http://vllab1.ucmerced.edu/~whung/adv-semi-seg/resnet101COCO-41f33a49.pth
# This script matches this pre-trained model's parameter dict keys with our implementation,
# We have 104 extra .num_batches_tracked, which are all tensor(0), others are the same
hung_coco_filename = 'resnet101COCO-41f33a49.pth'
coco = torch.load(hung_coco_filename)
voc_net = deeplab_v2(num_classes=21)
city_net = deeplab_v2(num_classes=19)
my_voc = voc_net.state_dict().copy()
my_city = city_net.state_dict().copy()
voc_shape_not_match = 0
voc_shape_match = 0
city_shape_not_match = 0
city_shape_match = 0

for key in coco:
    if 'layer5' in key:
        my_key = 'classifier.0.convs' + key.split('conv2d_list')[1]
    else:
        my_key = 'backbone.' + key
    if my_voc[my_key].shape == coco[key].shape:
        voc_shape_match += 1
        my_voc[my_key] = coco[key]
    else:
        voc_shape_not_match += 1