CNNCreator_Alexnet.py 13.2 KB
Newer Older
1
from caffe2.python import workspace, core, model_helper, brew, optimizer
2
3
4
5
from caffe2.python.predictor import mobile_exporter
from caffe2.proto import caffe2_pb2
import numpy as np

6
7
8
9
import logging
import os
import shutil
import sys
10
import cv2
11

12
13
#TODO: Check whether class is needed
#class CNNCreator_Alexnet:
14
15
16
17
18
19
20
21
22

module = None
_data_dir_ = "data/Alexnet/"
_model_dir_ = "model/Alexnet/"
_model_prefix_ = "Alexnet"
_input_names_ = ['data']
_input_shapes_ = [(3,224,224)]
_output_names_ = ['predictions_label']

23
24
EPOCHS     = 10000	# total training iterations
BATCH_SIZE = 256	# batch size for training
25
26
27
28
29
30
31
32
33
34
35
36
37
CONTEXT = 'gpu'
EVAL_METRIC = 'accuracy'
OPTIMIZER_TYPE = 'adam'
BASE_LEARNING_RATE = 0.001
WEIGHT_DECAY = 0.001
POLICY = 'fixed'
STEP_SIZE = 1
EPSILON = 1e-8
BETA1 = 0.9
BETA2 = 0.999
GAMMA = 0.999
MOMENTUM = 0.9

38
39
40
41
42

CURRENT_FOLDER      = os.path.join('./')
DATA_FOLDER         = os.path.join(CURRENT_FOLDER, 'data')
ROOT_FOLDER         = os.path.join(CURRENT_FOLDER, 'model')

43
44
45
46
#TODO: Modify paths to make them dynamic
#For Windows
#INIT_NET = 'D:/Yeverino/git_projects/Caffe2_scripts/caffe2_ema_cnncreator/init_net'
#PREDICT_NET = 'D:/Yeverino/git_projects/Caffe2_scripts/caffe2_ema_cnncreator/predict_net'
47

48
#For Ubuntu
49
50
INIT_NET = './model/init_net'
PREDICT_NET = './model/predict_net'
51

52
53
54
55
56
57
58
# Move into train function if test of deploy_net is removed
if CONTEXT == 'cpu':
    device_opts = core.DeviceOption(caffe2_pb2.CPU, 0)
    print("CPU mode selected")
elif CONTEXT == 'gpu':
    device_opts = core.DeviceOption(caffe2_pb2.CUDA, 0)
    print("GPU mode selected")
59

60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def add_input(model, batch_size, db, db_type, device_opts):
    with core.DeviceScope(device_opts):
        # load the data
        data_uint8, label = brew.db_input(
            model,
            blobs_out=["data_uint8", "label"],
            batch_size=batch_size,
            db=db,
            db_type=db_type,
        )
        # cast the data to float
        data = model.Cast(data_uint8, "data", to=core.DataType.FLOAT)

        # scale data from [0,255] down to [0,1]
        data = model.Scale(data, data, scale=float(1./256))

        # don't need the gradient for the backward pass
        data = model.StopGradient(data, data)
        return data, label

def create_model(model, data, device_opts):
81
82
	with core.DeviceScope(device_opts):

83
		data = data
84
85
		# data, output shape: {[3,224,224]}
  	
86
		conv1_ = brew.conv(model, data, 'conv1_', dim_in=1, dim_out=96, kernel=11, stride=4)
87
		# conv1_, output shape: {[96,55,55]}
88
89
90
91
92
93
        lrn1_ = mx.symbol.LRN(data=conv1_,
            alpha=0.0001,
            beta=0.75,
            knorm=2,
            nsize=5,
            name="lrn1_")
94
95
96
		pool1_ = brew.max_pool(model, lrn1_, 'pool1_', kernel=3, stride=2)
		# pool1_, output shape: {[96,27,27]}
		relu1_ = brew.relu(model, pool1_, pool1_)
97
98
99
100
        split1_ = mx.symbol.split(data=relu1_,
            num_outputs=2,
            axis=1,
            name="split1_")
101
		# split1_, output shape: {[48,27,27][48,27,27]}
102
        get2_1_ = split1_[0]
103
104
  		conv2_1_ = brew.conv(model, get2_1_, 'conv2_1_', dim_in=48, dim_out=128, kernel=5, stride=1)
		# conv2_1_, output shape: {[128,27,27]}
105
106
107
108
109
110
        lrn2_1_ = mx.symbol.LRN(data=conv2_1_,
            alpha=0.0001,
            beta=0.75,
            knorm=2,
            nsize=5,
            name="lrn2_1_")
111
112
113
		pool2_1_ = brew.max_pool(model, lrn2_1_, 'pool2_1_', kernel=3, stride=2)
		# pool2_1_, output shape: {[128,13,13]}
		relu2_1_ = brew.relu(model, pool2_1_, pool2_1_)
114
        get2_2_ = split1_[1]
115
116
  		conv2_2_ = brew.conv(model, get2_2_, 'conv2_2_', dim_in=48, dim_out=128, kernel=5, stride=1)
		# conv2_2_, output shape: {[128,27,27]}
117
118
119
120
121
122
        lrn2_2_ = mx.symbol.LRN(data=conv2_2_,
            alpha=0.0001,
            beta=0.75,
            knorm=2,
            nsize=5,
            name="lrn2_2_")
123
124
125
		pool2_2_ = brew.max_pool(model, lrn2_2_, 'pool2_2_', kernel=3, stride=2)
		# pool2_2_, output shape: {[128,13,13]}
		relu2_2_ = brew.relu(model, pool2_2_, pool2_2_)
126
127
128
        concatenate3_ = mx.symbol.concat(relu2_1_, relu2_2_,
            dim=1,
            name="concatenate3_")
129
130
131
132
		# concatenate3_, output shape: {[256,13,13]}
  		conv3_ = brew.conv(model, concatenate3_, 'conv3_', dim_in=256, dim_out=384, kernel=3, stride=1)
		# conv3_, output shape: {[384,13,13]}
		relu3_ = brew.relu(model, conv3_, conv3_)
133
134
135
136
        split3_ = mx.symbol.split(data=relu3_,
            num_outputs=2,
            axis=1,
            name="split3_")
137
		# split3_, output shape: {[192,13,13][192,13,13]}
138
        get4_1_ = split3_[0]
139
140
141
142
143
144
145
146
  		conv4_1_ = brew.conv(model, get4_1_, 'conv4_1_', dim_in=192, dim_out=192, kernel=3, stride=1)
		# conv4_1_, output shape: {[192,13,13]}
		relu4_1_ = brew.relu(model, conv4_1_, conv4_1_)
  		conv5_1_ = brew.conv(model, relu4_1_, 'conv5_1_', dim_in=192, dim_out=128, kernel=3, stride=1)
		# conv5_1_, output shape: {[128,13,13]}
		pool5_1_ = brew.max_pool(model, conv5_1_, 'pool5_1_', kernel=3, stride=2)
		# pool5_1_, output shape: {[128,6,6]}
		relu5_1_ = brew.relu(model, pool5_1_, pool5_1_)
147
        get4_2_ = split3_[1]
148
149
150
151
152
153
154
155
  		conv4_2_ = brew.conv(model, get4_2_, 'conv4_2_', dim_in=192, dim_out=192, kernel=3, stride=1)
		# conv4_2_, output shape: {[192,13,13]}
		relu4_2_ = brew.relu(model, conv4_2_, conv4_2_)
  		conv5_2_ = brew.conv(model, relu4_2_, 'conv5_2_', dim_in=192, dim_out=128, kernel=3, stride=1)
		# conv5_2_, output shape: {[128,13,13]}
		pool5_2_ = brew.max_pool(model, conv5_2_, 'pool5_2_', kernel=3, stride=2)
		# pool5_2_, output shape: {[128,6,6]}
		relu5_2_ = brew.relu(model, pool5_2_, pool5_2_)
156
157
158
        concatenate6_ = mx.symbol.concat(relu5_1_, relu5_2_,
            dim=1,
            name="concatenate6_")
159
160
161
162
		# concatenate6_, output shape: {[256,6,6]}
		fc6_ = brew.fc(model, concatenate6_, 'fc6_', dim_in=256 * 6 * 6, dim_out=4096)
		# fc6_, output shape: {[4096,1,1]}
		relu6_ = brew.relu(model, fc6_, fc6_)
163
164
165
        dropout6_ = mx.symbol.Dropout(data=relu6_,
            p=0.5,
            name="dropout6_")
166
167
168
		fc7_ = brew.fc(model, dropout6_, 'fc7_', dim_in=4096, dim_out=4096)
		# fc7_, output shape: {[4096,1,1]}
		relu7_ = brew.relu(model, fc7_, fc7_)
169
170
171
        dropout7_ = mx.symbol.Dropout(data=relu7_,
            p=0.5,
            name="dropout7_")
172
173
174
175
176
177
		fc8_ = brew.fc(model, dropout7_, 'fc8_', dim_in=4096, dim_out=10)
		# fc8_, output shape: {[10,1,1]}
		predictions = brew.softmax(model, fc8_, 'predictions')

		return predictions

178
# this adds the loss and optimizer
179
def add_training_operators(model, output, label, device_opts) :
180
	with core.DeviceScope(device_opts):
181
		xent = model.LabelCrossEntropy([output, label], 'xent')
182
183
184
		loss = model.AveragedLoss(xent, "loss")

		model.AddGradientOperators([loss])
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209

        if OPTIMIZER_TYPE == 'adam':
            if POLICY == 'step':
                opt = optimizer.build_adam(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, stepsize=STEP_SIZE, beta1=BETA1, beta2=BETA2, epsilon=EPSILON)
            elif POLICY == 'fixed' or POLICY == 'inv':
                opt = optimizer.build_adam(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, beta1=BETA1, beta2=BETA2, epsilon=EPSILON)
            print("adam optimizer selected")
        elif OPTIMIZER_TYPE == 'sgd':
            if POLICY == 'step':
                opt = optimizer.build_sgd(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, stepsize=STEP_SIZE, gamma=GAMMA, momentum=MOMENTUM)
            elif POLICY == 'fixed' or POLICY == 'inv':
                opt = optimizer.build_sgd(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, gamma=GAMMA, momentum=MOMENTUM)
            print("sgd optimizer selected")
        elif OPTIMIZER_TYPE == 'rmsprop':
            if POLICY == 'step':
                opt = optimizer.build_rms_prop(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, stepsize=STEP_SIZE, decay=GAMMA, momentum=MOMENTUM, epsilon=EPSILON)
            elif POLICY == 'fixed' or POLICY == 'inv':
                opt = optimizer.build_rms_prop(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, decay=GAMMA, momentum=MOMENTUM, epsilon=EPSILON)
            print("rmsprop optimizer selected")
        elif OPTIMIZER_TYPE == 'adagrad':
            if POLICY == 'step':
                opt = optimizer.build_adagrad(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, stepsize=STEP_SIZE, decay=GAMMA, epsilon=EPSILON)
            elif POLICY == 'fixed' or POLICY == 'inv':
                opt = optimizer.build_adagrad(model, base_learning_rate=BASE_LEARNING_RATE, policy=POLICY, decay=GAMMA, epsilon=EPSILON)
            print("adagrad optimizer selected")
210

211
212
def add_accuracy(model, output, label, device_opts):
    with core.DeviceScope(device_opts):
213
214
215
216
        if EVAL_METRIC == 'accuracy':
            accuracy = brew.accuracy(model, [output, label], "accuracy")
        elif EVAL_METRIC == 'top_k_accuracy':
            accuracy = brew.accuracy(model, [output, label], "accuracy", top_k=3)
217
218
        return accuracy

219
220
def train(INIT_NET, PREDICT_NET, epochs, batch_size, device_opts) :

221
222
223
224
225
226
227
228
229
	workspace.ResetWorkspace(ROOT_FOLDER)

	arg_scope = {"order": "NCHW"}
	# == Training model ==
	train_model= model_helper.ModelHelper(name="train_net", arg_scope=arg_scope)
	data, label = add_input(train_model, batch_size=batch_size, db=os.path.join(DATA_FOLDER, 'mnist-train-nchw-lmdb'), db_type='lmdb', device_opts=device_opts)
	predictions = create_model(train_model, data, device_opts=device_opts)
	add_training_operators(train_model, predictions, label, device_opts=device_opts)
	add_accuracy(train_model, predictions, label, device_opts)
230
	with core.DeviceScope(device_opts):
231
		brew.add_weight_decay(train_model, WEIGHT_DECAY)
232

233
	# Initialize and create the training network
234
	workspace.RunNetOnce(train_model.param_init_net)
235
	workspace.CreateNet(train_model.net, overwrite=True)
236

237
238
	# Main Training Loop
	print("== Starting Training for " + str(epochs) + " epochs ==")
239
	for j in range(0, epochs):
240
241
242
243
244
245
246
247
248
249
250
		workspace.RunNet(train_model.net)
		if j % 50 == 0:
			print 'Iter: ' + str(j) + ': ' + 'Loss ' + str(workspace.FetchBlob("loss")) + ' - ' + 'Accuracy ' + str(workspace.FetchBlob('accuracy'))
	print("Training done")

	print("== Running Test model ==")
	# == Testing model. ==
	test_model= model_helper.ModelHelper(name="test_net", arg_scope=arg_scope, init_params=False)
	data, label = add_input(test_model, batch_size=100, db=os.path.join(DATA_FOLDER, 'mnist-test-nchw-lmdb'), db_type='lmdb', device_opts=device_opts)
	predictions = create_model(test_model, data, device_opts=device_opts)
	add_accuracy(test_model, predictions, label, device_opts)
251
252
253
	workspace.RunNetOnce(test_model.param_init_net)
	workspace.CreateNet(test_model.net, overwrite=True)

254
255
256
257
258
259
260
261
262
263
	# Main Testing Loop
	# batch size:        100
	# iteration:         100
	# total test images: 10000
	test_accuracy = np.zeros(100)
	for i in range(100):
		# Run a forward pass of the net on the current batch
		workspace.RunNet(test_model.net)
		# Collect the batch accuracy from the workspace
		test_accuracy[i] = workspace.FetchBlob('accuracy')
264

265
	print('Test_accuracy: {:.4f}'.format(test_accuracy.mean()))
266

267
268
269
270
	# == Deployment model. ==
	# We simply need the main AddModel part.
	deploy_model = model_helper.ModelHelper(name="deploy_net", arg_scope=arg_scope, init_params=False)
	create_model(deploy_model, "data", device_opts)
271

272
	print("Saving deploy model")
273
	save_net(INIT_NET, PREDICT_NET, deploy_model)
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288

def save_net(init_net_path, predict_net_path, model):

	init_net, predict_net = mobile_exporter.Export(
		workspace,
		model.net,
		model.params
	)

	print("Save the model to init_net.pb and predict_net.pb")
	with open(predict_net_path + '.pb', 'wb') as f:
		f.write(model.net._net.SerializeToString())
	with open(init_net_path + '.pb', 'wb') as f:
		f.write(init_net.SerializeToString())

289
	print("Save the model to init_net.pbtxt and predict_net.pbtxt")
290
291
292
293
	with open(init_net_path + '.pbtxt', 'w') as f:
		f.write(str(init_net))
	with open(predict_net_path + '.pbtxt', 'w') as f:
		f.write(str(predict_net))
294
	print("== Saved init_net and predict_net ==")
295
296
297
298
299
300
301
302
303
304
305
306
307

def load_net(init_net_path, predict_net_path, device_opts):
	init_def = caffe2_pb2.NetDef()
	with open(init_net_path + '.pb', 'rb') as f:
		init_def.ParseFromString(f.read())
		init_def.device_option.CopyFrom(device_opts)
		workspace.RunNetOnce(init_def.SerializeToString())

	net_def = caffe2_pb2.NetDef()
	with open(predict_net_path + '.pb', 'rb') as f:
		net_def.ParseFromString(f.read())
		net_def.device_option.CopyFrom(device_opts)
		workspace.CreateNet(net_def.SerializeToString(), overwrite=True)
308
	print("== Loaded init_net and predict_net ==")
309

310
train(INIT_NET, PREDICT_NET, epochs=EPOCHS, batch_size=BATCH_SIZE, device_opts=device_opts)
311
312

print '\n********************************************'
313
print("Loading Deploy model")
314
315
load_net(INIT_NET, PREDICT_NET, device_opts=device_opts)

316
317
img = cv2.imread("3.jpg")                                   # Load test image
img = cv2.resize(img, (28,28))                              # Resize to 28x28
318
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY )               # Covert to grayscale
319
320
321
322
323
324
325
326
img = img.reshape((1,1,28,28)).astype('float32')            # Reshape to (1,1,28,28)
workspace.FeedBlob("data", img, device_option=device_opts)  # FeedBlob
workspace.RunNet('deploy_net', num_iter=1)                  # Forward

print("\nInput: {}".format(img.shape))
pred = workspace.FetchBlob("predictions") #TODO: Consider multiple output names
print("Output: {}".format(pred))
print("Output class: {}".format(np.argmax(pred)))