Upload CaffeLoader.py

CaffeLoader.py

@@ -0,0 +1,254 @@
+import torch
+import torch.nn as nn
+
+
+class VGG(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, num_classes),
+        )
+
+
+class VGG_SOD(nn.Module):
+    def __init__(self, features, num_classes=100):
+        super(VGG_SOD, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 100),
+        )
+
+
+class VGG_FCN32S(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG_FCN32S, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Conv2d(512,4096,(7, 7)),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.Conv2d(4096,4096,(1, 1)),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+        )
+
+
+class VGG_PRUNED(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG_PRUNED, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+        )
+
+
+class NIN(nn.Module):
+    def __init__(self, pooling):
+        super(NIN, self).__init__()
+        if pooling == 'max':
+            pool2d = nn.MaxPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
+        elif pooling == 'avg':
+            pool2d = nn.AvgPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
+
+        self.features = nn.Sequential(
+            nn.Conv2d(3,96,(11, 11),(4, 4)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(96,96,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(96,96,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Conv2d(96,256,(5, 5),(1, 1),(2, 2)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256,256,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256,256,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Conv2d(256,384,(3, 3),(1, 1),(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384,384,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384,384,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Dropout(0.5),
+            nn.Conv2d(384,1024,(3, 3),(1, 1),(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024,1024,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024,1000,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d((6, 6),(1, 1),(0, 0),ceil_mode=True),
+            nn.Softmax(),
+        )
+
+
+
+class ModelParallel(nn.Module):
+    def __init__(self, net, device_ids, device_splits):
+        super(ModelParallel, self).__init__()
+        self.device_list = self.name_devices(device_ids.split(','))
+        self.chunks = self.chunks_to_devices(self.split_net(net, device_splits.split(',')))
+
+    def name_devices(self, input_list):
+        device_list = []
+        for i, device in enumerate(input_list):
+            if str(device).lower() != 'c':
+                device_list.append("cuda:" + str(device))
+            else:
+                device_list.append("cpu")
+        return device_list
+
+    def split_net(self, net, device_splits):
+        chunks, cur_chunk = [], nn.Sequential()
+        for i, l in enumerate(net):
+            cur_chunk.add_module(str(i), net[i])
+            if str(i) in device_splits and device_splits != '':
+                del device_splits[0]
+                chunks.append(cur_chunk)
+                cur_chunk = nn.Sequential()
+        chunks.append(cur_chunk)
+        return chunks
+
+    def chunks_to_devices(self, chunks):
+        for i, chunk in enumerate(chunks):
+            chunk.to(self.device_list[i])
+        return chunks
+
+    def c(self, input, i):
+        if input.type() == 'torch.FloatTensor' and 'cuda' in self.device_list[i]:
+            input = input.type('torch.cuda.FloatTensor')
+        elif input.type() == 'torch.cuda.FloatTensor' and 'cpu' in self.device_list[i]:
+            input = input.type('torch.FloatTensor')
+        return input
+
+    def forward(self, input):
+        for i, chunk in enumerate(self.chunks):
+            if i < len(self.chunks) -1:
+                input = self.c(chunk(self.c(input, i).to(self.device_list[i])), i+1).to(self.device_list[i+1])
+            else:
+                input = chunk(input)
+        return input
+
+
+
+def buildSequential(channel_list, pooling):
+    layers = []
+    in_channels = 3
+    if pooling == 'max':
+        pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
+    elif pooling == 'avg':
+        pool2d = nn.AvgPool2d(kernel_size=2, stride=2)
+    else:
+        raise ValueError("Unrecognized pooling parameter")
+    for c in channel_list:
+        if c == 'P':
+            layers += [pool2d]
+        else:
+            conv2d = nn.Conv2d(in_channels, c, kernel_size=3, padding=1)
+            layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = c
+    return nn.Sequential(*layers)
+
+
+channel_list = {
+'VGG-16p': [24, 22, 'P', 41, 51, 'P', 108, 89, 111, 'P', 184, 276, 228, 'P', 512, 512, 512, 'P'],
+'VGG-16': [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 'P', 512, 512, 512, 'P', 512, 512, 512, 'P'],
+'VGG-19': [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 256, 'P', 512, 512, 512, 512, 'P', 512, 512, 512, 512, 'P'],
+}
+
+nin_dict = {
+'C': ['conv1', 'cccp1', 'cccp2', 'conv2', 'cccp3', 'cccp4', 'conv3', 'cccp5', 'cccp6', 'conv4-1024', 'cccp7-1024', 'cccp8-1024'],
+'R': ['relu0', 'relu1', 'relu2', 'relu3', 'relu5', 'relu6', 'relu7', 'relu8', 'relu9', 'relu10', 'relu11', 'relu12'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4'],
+'D': ['drop'],
+}
+vgg16_dict = {
+'C': ['conv1_1', 'conv1_2', 'conv2_1', 'conv2_2', 'conv3_1', 'conv3_2', 'conv3_3', 'conv4_1', 'conv4_2', 'conv4_3', 'conv5_1', 'conv5_2', 'conv5_3'],
+'R': ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1', 'relu3_2', 'relu3_3', 'relu4_1', 'relu4_2', 'relu4_3', 'relu5_1', 'relu5_2', 'relu5_3'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4', 'pool5'],
+}
+vgg19_dict = {
+'C': ['conv1_1', 'conv1_2', 'conv2_1', 'conv2_2', 'conv3_1', 'conv3_2', 'conv3_3', 'conv3_4', 'conv4_1', 'conv4_2', 'conv4_3', 'conv4_4', 'conv5_1', 'conv5_2', 'conv5_3', 'conv5_4'],
+'R': ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1', 'relu3_2', 'relu3_3', 'relu3_4', 'relu4_1', 'relu4_2', 'relu4_3', 'relu4_4', 'relu5_1', 'relu5_2', 'relu5_3', 'relu5_4'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4', 'pool5'],
+}
+
+
+def modelSelector(model_file, pooling):
+    vgg_list = ["fcn32s", "pruning", "sod", "vgg"]
+    if any(name in model_file for name in vgg_list):
+        if "pruning" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using The Channel Pruning Model")
+            cnn, layerList = VGG_PRUNED(buildSequential(channel_list['VGG-16p'], pooling)), vgg16_dict
+        elif "fcn32s" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using the fcn32s-heavy-pascal Model")
+            cnn, layerList = VGG_FCN32S(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        elif "sod" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using The SOD Fintune Model")
+            cnn, layerList = VGG_SOD(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        elif "19" in model_file:
+            print("VGG-19 Architecture Detected")
+            cnn, layerList = VGG(buildSequential(channel_list['VGG-19'], pooling)), vgg19_dict
+        elif "16" in model_file:
+            print("VGG-16 Architecture Detected")
+            cnn, layerList = VGG(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        else:
+            raise ValueError("VGG architecture not recognized.")
+    elif "nin" in model_file:
+        print("NIN Architecture Detected")
+        cnn, layerList = NIN(pooling), nin_dict
+    else:
+        raise ValueError("Model architecture not recognized.")
+    return cnn, layerList
+
+
+# Print like Torch7/loadcaffe
+def print_loadcaffe(cnn, layerList):
+    c = 0
+    for l in list(cnn):
+         if "Conv2d" in str(l):
+             in_c, out_c, ks  = str(l.in_channels), str(l.out_channels), str(l.kernel_size)
+             print(layerList['C'][c] +": " +  (out_c + " " + in_c + " " + ks).replace(")",'').replace("(",'').replace(",",'') )
+             c+=1
+         if c == len(layerList['C']):
+             break
+
+
+# Load the model, and configure pooling layer type
+def loadCaffemodel(model_file, pooling, use_gpu, disable_check):
+    cnn, layerList = modelSelector(str(model_file).lower(), pooling)
+
+    cnn.load_state_dict(torch.load(model_file), strict=(not disable_check))
+    print("Successfully loaded " + str(model_file))
+
+    # Maybe convert the model to cuda now, to avoid later issues
+    if "c" not in str(use_gpu).lower() or "c" not in str(use_gpu[0]).lower():
+        cnn = cnn.cuda()
+    cnn = cnn.features
+
+    print_loadcaffe(cnn, layerList)
+
+    return cnn, layerList