Make Transfer learning of SqueezeNet on Caffe2

Kazami
2 min readJul 31, 2017

--

In modern computer vision applicaiton, seldom company train the neural network from scratch, because of lacking of big size data and computing resource. In this article, we would focus on how to make transfer learning on Caffe2.

Before operation, it should load two protobuf files which are pred_net.pb and init_net.pb.

First, we download the two pretrained pb file (model zoo). wThese models are trained from the enourmous company, such as Google, Facebook and Microsoft. For example, in the squeezenet model, we remove the final conv and softmax layer. In order to remove the final conv layer, the protobuf file is parsed and splitted to the fixed feature layers.

def splitmodel(prednet, initnet, target):
lookup = {}
forward = False
for layer in predict_net.op:
for inop in layer.input:
tmp = ""
for out_op in layer.output:
if forward:
tmp = out if forward else inop
if forward:
if inop in lookup:
lookup[inop].append(tmp)
else:
lookup[inop] = []
lookup[inop].append(tmp)
else:
if outop in lookup:
lookup[outop].append(tmp)
else:
lookup[outop] = []
lookup[outop].append(tmp)
pred = caffe2_pb2.NetDef()
init = caffe2_pb2.NetDef()
with open(PREDICT_NET) as f:
pred.ParseFromString(f.read())
with open(INIT_NET) as g:
init.ParseFromString(g.read())
import os
MODEL = 'squeezenet', 'init_net.pb', 'predict_net.pb'
#init_net.pb is in the ~/squeezenet/init_net.pb
#predict_net.pb is in the ~/squeezenet/predict_net.pb
CAFFE_MODELS = "/home/kazami/"
INIT_NET = os.path.join(CAFFE_MODELS, MODEL[0], MODEL[1])
PREDICT_NET = os.path.join(CAFFE_MODELS, MODEL[0], MODEL[2])
import dataset.feature as feature
init_n, pred_n = feature.splitmodel(pred, init, 'pool10')
pren = workspace.Predictor(pred_n.SerializeToString(), init_n.SerializeToString())

After acquring the fixed layers, the Predictor helper function could be used to extract features. These features could be stored in the leveldb or minidb as a features db, which is similar to making training dataset. The only difference is that are features rather than raw images.

import dataset.imgfolder as folder
import dataset.imgloader as loader
import db.dbwriter as dbw
import numpy as np
def folder_load(foldername):
dataset = folder.ImageFolder(foldername)
features = []
labels = []
d = dbw.DB('minidb', 'featureminidb')

imgs = loader.ImgLoader(dataset)
for i, data in enumerate(imgs):
img, lab = data
img = img[np.newaxis, :, :, :].astype(np.float32)
feature = pre.run([img])
print (feature[0][0].shape)
features.append(feature[0][0])
labels.append(lab)
if len(features) == 50:
features = np.array(features)
labels = np.array(labels)
d.writedb(features, labels)
features = []
labels = []

Then, we would create new layer and re-train new net.

The above code is as follows:

Ref:

  1. http://cs231n.github.io/transfer-learning/

--

--