来源地址:https://caffe2.ai/docs/tutorial-MNIST.html
1.导入necessities
%matplotlib inline from matplotlib import pyplot import numpy as np import os import shutil import caffe2.python.predictor.predictor_exporter as pe from caffe2.python import core, model_helper, net_drawer, workspace, visualize, brew # If you would like to see some really detailed initializations, # you can change --caffe2_log_level=0 to --caffe2_log_level=-1 core.GlobalInit(['caffe2', '--caffe2_log_level=0']) print("Necessities imported!")
2.设置训练集、测试集
# This section preps your image and test set in a lmdb database def DownloadResource(url, path): '''Downloads resources from s3 by url and unzips them to the provided path''' import requests, zipfile, StringIO print("Downloading... {} to {}".format(url, path)) r = requests.get(url, stream=True) z = zipfile.ZipFile(StringIO.StringIO(r.content)) z.extractall(path) print("Completed download and extraction.") current_folder = os.path.join(os.path.expanduser('~'), 'caffe2_notebooks') data_folder = os.path.join(current_folder, 'tutorial_data', 'mnist') root_folder = os.path.join(current_folder, 'tutorial_files', 'tutorial_mnist') db_missing = False if not os.path.exists(data_folder): os.makedirs(data_folder) print("Your data folder was not found!! This was generated: {}".format(data_folder)) # Look for existing database: lmdb if os.path.exists(os.path.join(data_folder,"mnist-train-nchw-lmdb")): print("lmdb train db found!") else: db_missing = True if os.path.exists(os.path.join(data_folder,"mnist-test-nchw-lmdb")): print("lmdb test db found!") else: db_missing = True # attempt the download of the db if either was missing if db_missing: print("one or both of the MNIST lmbd dbs not found!!") db_url = "http://download.caffe2.ai/databases/mnist-lmdb.zip" try: DownloadResource(db_url, data_folder) except Exception as ex: print("Failed to download dataset. Please download it manually from {}".format(db_url)) print("Unzip it and place the two database folders here: {}".format(data_folder)) raise ex if os.path.exists(root_folder): print("Looks like you ran this before, so we need to cleanup those old files...") shutil.rmtree(root_folder) os.makedirs(root_folder) workspace.ResetWorkspace(root_folder) print("training data folder:" + data_folder) print("workspace root folder:" + root_folder)
四个模块
(1) The data input part (AddInput function)
(2) The main computation part (AddLeNetModel function)
(3) The training part - adding gradient operators, update, etc. (AddTrainingOperators function)
(4) The bookkeeping part, where we just print out statistics for inspection. (AddBookkeepingOperators function)
3.从数据库加载数据
[batch_size, 1, 28, 28]
数据类型为uint8,标签为int,将数据从[0,255]转为[0,1]
def AddInput(model, batch_size, db, db_type): # load the data data_uint8, label = model.TensorProtosDBInput( [], ["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
4.添加LeNet模型
dim_in是输入通道数,dim_out是输出通道数
def AddLeNetModel(model, data): ''' This part is the standard LeNet model: from data to the softmax prediction. For each convolutional layer we specify dim_in - number of input channels and dim_out - number or output channels. Also each Conv and MaxPool layer changes the image size. For example, kernel of size 5 reduces each side of an image by 4. While when we have kernel and stride sizes equal 2 in a MaxPool layer, it divides each side in half. ''' # Image size: 28 x 28 -> 24 x 24 conv1 = brew.conv(model, data, 'conv1', dim_in=1, dim_out=20, kernel=5) # Image size: 24 x 24 -> 12 x 12 pool1 = brew.max_pool(model, conv1, 'pool1', kernel=2, stride=2) # Image size: 12 x 12 -> 8 x 8 conv2 = brew.conv(model, pool1, 'conv2', dim_in=20, dim_out=100, kernel=5) # Image size: 8 x 8 -> 4 x 4 pool2 = brew.max_pool(model, conv2, 'pool2', kernel=2, stride=2) # 50 * 4 * 4 stands for dim_out from previous layer multiplied by the image size fc3 = brew.fc(model, pool2, 'fc3', dim_in=100 * 4 * 4, dim_out=500) #fc3 = brew.relu(model, fc3, fc3) relu = brew.relu(model, fc3, fc3) pred = brew.fc(model, relu, 'pred', 500, 10) softmax = brew.softmax(model, pred, 'softmax') return softmax添加精度计算
def AddAccuracy(model, softmax, label): """Adds an accuracy op to the model""" accuracy = brew.accuracy(model, [softmax, label], "accuracy") return accuracy5.添加训练操作
def AddTrainingOperators(model, softmax, label): """Adds training operators to the model.""" xent = model.LabelCrossEntropy([softmax, label], 'xent') # compute the expected loss loss = model.AveragedLoss(xent, "loss") # track the accuracy of the model AddAccuracy(model, softmax, label) # use the average loss we just computed to add gradient operators to the model model.AddGradientOperators([loss]) # do a simple stochastic gradient descent ITER = brew.iter(model, "iter") # set the learning rate schedule LR = model.LearningRate( ITER, "LR", base_lr=-0.1, policy="step", stepsize=1, gamma=0.999 ) # ONE is a constant value that is used in the gradient update. We only need # to create it once, so it is explicitly placed in param_init_net. ONE = model.param_init_net.ConstantFill([], "ONE", shape=[1], value=1.0) # Now, for each parameter, we do the gradient updates. for param in model.params: # Note how we get the gradient of each parameter - ModelHelper keeps # track of that. param_grad = model.param_to_grad[param] # The update is a simple weighted sum: param = param + param_grad * LR model.WeightedSum([param, ONE, param_grad, LR], param)
6.分类登记(不影响训练过程)
def AddBookkeepingOperators(model): """This adds a few bookkeeping operators that we can inspect later. These operators do not affect the training procedure: they only collect statistics and prints them to file or to logs. """ # Print basically prints out the content of the blob. to_file=1 routes the # printed output to a file. The file is going to be stored under # root_folder/[blob name] model.Print('accuracy', [], to_file=1) model.Print('loss', [], to_file=1) # Summarizes the parameters. Different from Print, Summarize gives some # statistics of the parameter, such as mean, std, min and max. for param in model.params: model.Summarize(param, [], to_file=1) model.Summarize(model.param_to_grad[param], [], to_file=1) # Now, if we really want to be verbose, we can summarize EVERY blob # that the model produces; it is probably not a good idea, because that # is going to take time - summarization do not come for free. For this # demo, we will only show how to summarize the parameters and their # gradients.
7.NCHW存储顺序
arg_scope = {"order": "NCHW"} train_model = model_helper.ModelHelper(name="mnist_train", arg_scope=arg_scope) data, label = AddInput( train_model, batch_size=64, db=os.path.join(data_folder, 'mnist-train-nchw-lmdb'), db_type='lmdb') softmax = AddLeNetModel(train_model, data) AddTrainingOperators(train_model, softmax, label) AddBookkeepingOperators(train_model) # Testing model. We will set the batch size to 100, so that the testing # pass is 100 iterations (10,000 images in total). # For the testing model, we need the data input part, the main LeNetModel # part, and an accuracy part. Note that init_params is set False because # we will be using the parameters obtained from the train model. test_model = model_helper.ModelHelper( name="mnist_test", arg_scope=arg_scope, init_params=False) data, label = AddInput( test_model, batch_size=100, db=os.path.join(data_folder, 'mnist-test-nchw-lmdb'), db_type='lmdb') softmax = AddLeNetModel(test_model, data) AddAccuracy(test_model, softmax, label) # Deployment model. We simply need the main LeNetModel part. deploy_model = model_helper.ModelHelper( name="mnist_deploy", arg_scope=arg_scope, init_params=False) AddLeNetModel(deploy_model, "data") # You may wonder what happens with the param_init_net part of the deploy_model. # No, we will not use them, since during deployment time we will not randomly # initialize the parameters, but load the parameters from the db.
8.观看训练与部署模型的外观(需要conda install graphviz)
from IPython import display graph = net_drawer.GetPydotGraph(train_model.net.Proto().op, "mnist", rankdir="LR") display.Image(graph.create_png(), width=800)
graph = net_drawer.GetPydotGraphMinimal( train_model.net.Proto().op, "mnist", rankdir="LR", minimal_dependency=True) display.Image(graph.create_png(), width=800)
print(str(train_model.param_init_net.Proto())[:400] + '\n...')
9.将协议缓冲区转储到磁盘
with open(os.path.join(root_folder, "train_net.pbtxt"), 'w') as fid: fid.write(str(train_model.net.Proto())) with open(os.path.join(root_folder, "train_init_net.pbtxt"), 'w') as fid: fid.write(str(train_model.param_init_net.Proto())) with open(os.path.join(root_folder, "test_net.pbtxt"), 'w') as fid: fid.write(str(test_model.net.Proto())) with open(os.path.join(root_folder, "test_init_net.pbtxt"), 'w') as fid: fid.write(str(test_model.param_init_net.Proto())) with open(os.path.join(root_folder, "deploy_net.pbtxt"), 'w') as fid: fid.write(str(deploy_model.net.Proto())) print("Protocol buffers files have been created in your root folder: " + root_folder)
10.运行训练程序
# The parameter initialization network only needs to be run once. workspace.RunNetOnce(train_model.param_init_net) # creating the network workspace.CreateNet(train_model.net, overwrite=True) # set the number of iterations and track the accuracy & loss total_iters = 200 accuracy = np.zeros(total_iters) loss = np.zeros(total_iters) # Now, we will manually run the network for 200 iterations. for i in range(total_iters): workspace.RunNet(train_model.net) accuracy[i] = workspace.FetchBlob('accuracy') loss[i] = workspace.FetchBlob('loss') # After the execution is done, let's plot the values. pyplot.plot(loss, 'b') pyplot.plot(accuracy, 'r') pyplot.legend(('Loss', 'Accuracy'), loc='upper right'
11.对数据和预测进行抽样
# Let's look at some of the data. pyplot.figure() data = workspace.FetchBlob('data') _ = visualize.NCHW.ShowMultiple(data) pyplot.figure() softmax = workspace.FetchBlob('softmax') _ = pyplot.plot(softmax[0], 'ro') pyplot.title('Prediction for the first image')
# Convolutions for this mini-batch pyplot.figure() conv = workspace.FetchBlob('conv1') shape = list(conv.shape) shape[1] = 1 # We can look into any channel. This of it as a feature model learned conv = conv[:,15,:,:].reshape(shape) _ = visualize.NCHW.ShowMultiple(conv)
12.运行测试报告测试精度
# run a test pass on the test net workspace.RunNetOnce(test_model.param_init_net) workspace.CreateNet(test_model.net, overwrite=True) test_accuracy = np.zeros(100) for i in range(100): workspace.RunNet(test_model.net.Proto().name) test_accuracy[i] = workspace.FetchBlob('accuracy') # After the execution is done, let's plot the values. pyplot.plot(test_accuracy, 'r') pyplot.title('Acuracy over test batches.') print('test_accuracy: %f' % test_accuracy.mean())
13.保存模型(包括训练权重和偏差)
# construct the model to be exported # the inputs/outputs of the model are manually specified. pe_meta = pe.PredictorExportMeta( predict_net=deploy_model.net.Proto(), parameters=[str(b) for b in deploy_model.params], inputs=["data"], outputs=["softmax"], ) # save the model to a file. Use minidb as the file format pe.save_to_db("minidb", os.path.join(root_folder, "mnist_model.minidb"), pe_meta) print("The deploy model is saved to: " + root_folder + "/mnist_model.minidb")
14.加载模型与验证
# we retrieve the last input data out and use it in our prediction test before we scratch the workspace blob = workspace.FetchBlob("data") pyplot.figure() _ = visualize.NCHW.ShowMultiple(blob) # reset the workspace, to make sure the model is actually loaded workspace.ResetWorkspace(root_folder) # verify that all blobs are destroyed. print("The blobs in the workspace after reset: {}".format(workspace.Blobs())) # load the predict net predict_net = pe.prepare_prediction_net(os.path.join(root_folder, "mnist_model.minidb"), "minidb") # verify that blobs are loaded back print("The blobs in the workspace after loading the model: {}".format(workspace.Blobs())) # feed the previously saved data to the loaded model workspace.FeedBlob("data", blob) # predict workspace.RunNetOnce(predict_net) softmax = workspace.FetchBlob("softmax") # the first letter should be predicted correctly pyplot.figure() _ = pyplot.plot(softmax[0], 'ro') pyplot.title('Prediction for the first image')