def __init__(self, decoder, output_size, in_channels=3, pretrained=True): super(MobileNet, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained: pretrained_path = os.path.join('imagenet', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') checkpoint = torch.load(pretrained_path) state_dict = checkpoint['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) if in_channels == 3: self.mobilenet = nn.Sequential(*(mobilenet.model[i] for i in range(14))) else: def conv_bn(inp, oup, stride): return nn.Sequential( nn.Conv2d(inp, oup, 3, stride, 1, bias=False), nn.BatchNorm2d(oup), nn.ReLU6(inplace=True) ) self.mobilenet = nn.Sequential( conv_bn(in_channels, 32, 2), *(mobilenet.model[i] for i in range(1,14)) ) self.decoder = choose_decoder(decoder)
def __init__(self, output_size, pretrained=True): super(MobileNetSkipConcat, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained: pretrained_path = os.path.join('imagenet', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') checkpoint = torch.load(pretrained_path) state_dict = checkpoint['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr( self, 'conv{}'.format(i), mobilenet.model[i]) kernel_size = 5 # self.decode_conv1 = conv(1024, 512, kernel_size) # self.decode_conv2 = conv(512, 256, kernel_size) # self.decode_conv3 = conv(256, 128, kernel_size) # self.decode_conv4 = conv(128, 64, kernel_size) # self.decode_conv5 = conv(64, 32, kernel_size) self.decode_conv1 = nn.Sequential( depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential( depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential( depthwise(512, kernel_size), pointwise(512, 128)) self.decode_conv4 = nn.Sequential( depthwise(256, kernel_size), pointwise(256, 64)) self.decode_conv5 = nn.Sequential( depthwise(128, kernel_size), pointwise(128, 32)) self.upconstant = [[512,8,8],[256,16,16],[128,32,32],[64,64,64],[32,128,128]] self.decode_conv6 = nn.Sequential( nn.Conv2d(32+3,3,1,1,0,bias=False), nn.BatchNorm2d(3), nn.Sigmoid(), ) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)
def __init__(self, output_size, pretrained_encoder=True): super(MobileNetSkipAdd, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained_encoder: pretrained_path = os.path.join( os.getcwd(), "imagenet", "results", "imagenet.arch=mobilenet.lr=0.1.bs=256", "model_best.pth.tar", ) checkpoint = torch.load(pretrained_path, map_location=torch.device("cpu")) state_dict = checkpoint["state_dict"] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr(self, "conv{}".format(i), mobilenet.model[i]) kernel_size = 5 # self.decode_conv1 = conv(1024, 512, kernel_size) # self.decode_conv2 = conv(512, 256, kernel_size) # self.decode_conv3 = conv(256, 128, kernel_size) # self.decode_conv4 = conv(128, 64, kernel_size) # self.decode_conv5 = conv(64, 32, kernel_size) self.decode_conv1 = nn.Sequential(depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential(depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential(depthwise(256, kernel_size), pointwise(256, 128)) self.decode_conv4 = nn.Sequential(depthwise(128, kernel_size), pointwise(128, 64)) self.decode_conv5 = nn.Sequential(depthwise(64, kernel_size), pointwise(64, 32)) self.decode_conv6 = pointwise(32, 1) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)
def __init__(self, output_size, pretrained=True): super(MobileNetSkipAdd, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained: # pretrained_path = os.path.join('imagenet', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') # pretrained_path = os.path.join('/content/drive/My Drive/FastDepth', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') pretrained_path = os.path.join('mobilnet', 'model_best.pth.tar') print("Pre Trained Mobile net path : " + pretrained_path) checkpoint = torch.load(pretrained_path) state_dict = checkpoint['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr(self, 'conv{}'.format(i), mobilenet.model[i]) kernel_size = 5 # self.decode_conv1 = conv(1024, 512, kernel_size) # self.decode_conv2 = conv(512, 256, kernel_size) # self.decode_conv3 = conv(256, 128, kernel_size) # self.decode_conv4 = conv(128, 64, kernel_size) # self.decode_conv5 = conv(64, 32, kernel_size) self.decode_conv1 = nn.Sequential(depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential(depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential(depthwise(256, kernel_size), pointwise(256, 128)) self.decode_conv4 = nn.Sequential(depthwise(128, kernel_size), pointwise(128, 64)) self.decode_conv5 = nn.Sequential(depthwise(64, kernel_size), pointwise(64, 32)) self.decode_conv6 = pointwise(32, 1) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)
def __init__(self, output_size, pretrained=True): super(MobileNetSkipAdd, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained: pretrained_path = os.path.join( '..', 'models', 'imagenet_pretrained_mobilenet.pth.tar') checkpoint = torch.load(pretrained_path, map_location='cpu') state_dict = checkpoint['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr(self, 'conv{}'.format(i), mobilenet.model[i]) # TODO: separate to encoder and decoder? the gaussian prior thing, add before this somehow? kernel_size = 5 # self.decode_conv1 = conv(1024, 512, kernel_size) # self.decode_conv2 = conv(512, 256, kernel_size) # self.decode_conv3 = conv(256, 128, kernel_size) # self.decode_conv4 = conv(128, 64, kernel_size) # self.decode_conv5 = conv(64, 32, kernel_size) self.decode_conv1 = nn.Sequential(depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential(depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential(depthwise(256, kernel_size), pointwise(256, 128)) self.decode_conv4 = nn.Sequential(depthwise(128, kernel_size), pointwise(128, 64)) self.decode_conv5 = nn.Sequential(depthwise(64, kernel_size), pointwise(64, 32)) self.decode_conv6 = pointwise(32, 1) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)
def __init__(self, output_size, pretrained=True): super(CustomMobilenetSemseg, self).__init__() self.output_size = output_size mobilenet = imagenet.mobilenet.MobileNet() if pretrained: pretrained_path = os.path.join( 'imagenet', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') checkpoint = torch.load(pretrained_path) state_dict = checkpoint['state_dict'] new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr(self, 'conv{}'.format(i), mobilenet.model[i]) kernel_size = 5 self.decode_conv1 = nn.Sequential(depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential(depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential(depthwise(256, kernel_size), pointwise(256, 128)) self.decode_conv4 = nn.Sequential(depthwise(128, kernel_size), pointwise(128, 64)) self.decode_conv5 = nn.Sequential(depthwise(64, kernel_size), pointwise(64, 32)) self.decode_conv6 = pointwise(32, 13) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)
def __init__(self, output_size, decoder='nnconv5dw', pretrained=True): super(MobileNetSkipConcat, self).__init__() assert (decoder in ['nnconv5dw','blconv5dw','upproj','shuffle','upconv','deconv']) self.output_size = output_size self.decoder = decoder mobilenet = imagenet.mobilenet.MobileNet() if pretrained: pretrained_path = os.path.join('imagenet', 'results', 'imagenet.arch=mobilenet.lr=0.1.bs=256', 'model_best.pth.tar') checkpoint = torch.load(pretrained_path) state_dict = checkpoint['state_dict'] from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # remove `module.` new_state_dict[name] = v mobilenet.load_state_dict(new_state_dict) else: mobilenet.apply(weights_init) for i in range(14): setattr( self, 'conv{}'.format(i), mobilenet.model[i]) kernel_size = 5 #self.decode_conv1 = conv(1024, 512, kernel_size) # self.decode_conv2 = conv(512, 256, kernel_size) # self.decode_conv3 = conv(256, 128, kernel_size) # self.decode_conv4 = conv(128, 64, kernel_size) # self.decode_conv5 = conv(64, 32, kernel_size) upsample = None if decoder == 'upproj': upsample = upproj elif decoder == 'upconv': upsample = upconv elif decoder == 'shuffle': upsample = shuffle_conv elif decoder == 'deconv': upsample = deconv if decoder in ['nnconv5dw','blconv5dw']: self.decode_conv1 = nn.Sequential( depthwise(1024, kernel_size), pointwise(1024, 512)) self.decode_conv2 = nn.Sequential( depthwise(512, kernel_size), pointwise(512, 256)) self.decode_conv3 = nn.Sequential( depthwise(512, kernel_size), pointwise(512, 128)) self.decode_conv4 = nn.Sequential( depthwise(256, kernel_size), pointwise(256, 64)) self.decode_conv5 = nn.Sequential( # Reduce channels depthwise(128, kernel_size), pointwise(128, 32)) else: self.decode_conv1 = upsample(1024,512) self.decode_conv2 = upsample(512,256) self.decode_conv3 = upsample(512,128) # Concat => inp = 256*2 self.decode_conv4 = upsample(256,64) # inp = 128*2 self.decode_conv5 = upsample(128,32) # inp = 64*2 self.decode_conv6 = pointwise(32, 1) weights_init(self.decode_conv1) weights_init(self.decode_conv2) weights_init(self.decode_conv3) weights_init(self.decode_conv4) weights_init(self.decode_conv5) weights_init(self.decode_conv6)