def __init__(self, model_name, precision, image_shape, batch_size):
self.float_type = np.float32 if precision == 'fp32' else np.float16
self.input = np.random.rand(batch_size, 3, image_shape[0],
image_shape[1]).astype(self.float_type)
with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
model = model_helper.ModelHelper(name=model_name, init_params=True)
softmax = getattr(self, model_name + '_model')(model, "data",
precision)
if precision == 'fp16':
softmax = model.net.HalfToFloat(softmax, softmax + "_fp32")
forward_net = core.Net(model.net.Proto())
if False:
print(forward_net.Proto())
from caffe2.python import net_drawer
graph = net_drawer.GetPydotGraphMinimal(forward_net,
rankdir="LR")
with open('/host/graph.png', 'wb') as fout:
fout.write(graph.create_png())
# bogus loss function to force backprop
loss = model.net.Sum(softmax, 'loss')
model.AddGradientOperators([loss])
workspace.RunNetOnce(model.param_init_net)
data = np.zeros((batch_size, 3, image_shape[0], image_shape[1]),
dtype=self.float_type)
workspace.FeedBlob("data", data)
workspace.CreateNet(model.net)
workspace.CreateNet(forward_net)
self.model = model
self.forward_net = forward_net
def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model,
cfg.TEST.WEIGHTS,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
graph = net_drawer.GetPydotGraphMinimal(model.net.Proto().op,
rankdir="LR",
minimal_dependency=True)
#display.Image(graph.create_png(), height=1800)
graph.write_png('graph.png')
with open(os.path.join("video/keypoint", "net.pbtxt"), 'w') as fid:
fid.write(str(model.net.Proto()))
return model
def save_graph_base(net, file_name, graph_name="net", op_only=True, blob_rename_func=None):
graph = None
ops = net.op
if blob_rename_func is not None:
ops = _modify_blob_names(ops, blob_rename_func)
if not op_only:
graph = net_drawer.GetPydotGraph(ops, graph_name, rankdir="TB")
else:
graph = net_drawer.GetPydotGraphMinimal(
ops, graph_name, rankdir="TB", minimal_dependency=True
)
try:
par_dir = os.path.dirname(file_name)
if not os.path.exists(par_dir):
os.makedirs(par_dir)
format = os.path.splitext(os.path.basename(file_name))[-1]
if format == ".png":
graph.write_png(file_name)
elif format == ".pdf":
graph.write_pdf(file_name)
elif format == ".svg":
graph.write_svg(file_name)
else:
print("Incorrect format {}".format(format))
except Exception as e:
print("Error when writing graph to image {}".format(e))
return graph
def save_graph(net, file_name, graph_name="net", op_only=True):
from caffe2.python import net_drawer
graph = None
ops = net.op
if not op_only:
graph = net_drawer.GetPydotGraph(ops, graph_name, rankdir="TB")
else:
graph = net_drawer.GetPydotGraphMinimal(ops,
graph_name,
rankdir="TB",
minimal_dependency=True)
try:
graph.write_png(file_name)
except Exception as e:
print('Error when writing graph to image {}'.format(e))
def visualize_net(model, path, minimal=False):
from caffe2.python import net_drawer
net_name = model.net.Proto().name
if minimal:
graph = net_drawer.GetPydotGraphMinimal(model.net.Proto().op,
net_name,
minimal_dependency=True)
else:
graph = net_drawer.GetPydotGraph(
model.net.Proto().op,
net_name,
)
if minimal:
img_output_path = os.path.join(path, net_name + '_minimal.png')
else:
img_output_path = os.path.join(path, net_name + '_full.png')
with open(img_output_path, 'w') as fopen:
fopen.write(graph.create_png())
logger.info("{}: Net image saved to: {}".format(
net_name, img_output_path))
def save_deploy_model(device_opts):
# save net forward only !!
deploy_model = model_helper.ModelHelper(name="deploy_net",
init_params=False)
last_out = create_resnet(model=deploy_model,
data='data',
num_input_channels=3,
num_groups=args.depths,
num_labels=10,
device_opts=device_opts,
is_test=True)
add_sortmax(deploy_model, last_out, device_opts)
workspace.RunNetOnce(deploy_model.param_init_net)
workspace.CreateNet(deploy_model.net, overwrite=True)
# save network images
graph = net_drawer.GetPydotGraphMinimal(deploy_model)
graph.write_svg('net.svg')
save_net(args.init_net, args.predict_net, deploy_model)
def generate_network_graph(model, config, tag="", use_mini=True):
''' generate the svg graph that show the network structure
Args:
use_mini: whether show full nodes of blobs and operators
tag: user specific name-tag
'''
# resolving necessary params
graph_dir = os.path.join(config['root_dir'], "output")
if tag != "":
tag = tag + "_"
# draw graph
if not use_mini:
# full-graph
full_graph = net_drawer.GetPydotGraph(
model.net.Proto().op,
"full_graph",
rankdir="TB",
)
graph_path = os.path.join(
graph_dir,
"{}_{}_{}graph.svg".format(config['model_name'],
config['dataset_name'], tag),
)
full_graph.write_svg(graph_path)
else:
# mini-graph
mini_graph = net_drawer.GetPydotGraphMinimal(model.net.Proto().op,
"mini_graph",
rankdir="TB",
minimal_dependency=True)
graph_path = os.path.join(
graph_dir,
"{}_{}_{}graph_minimal.svg".format(config['model_name'],
config['dataset_name'], tag),
)
mini_graph.write_svg(graph_path)
print("[INFO] write graph over...")
# 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.
graph = net_drawer.GetPydotGraph(train_model.net.Proto().op,
"mnist",
rankdir="LR")
graph.write_svg('Second.svg')
graph = net_drawer.GetPydotGraphMinimal(train_model.net.Proto().op,
"mnist",
rankdir="LR",
minimal_dependency=True)
graph.write_svg('Third.svg')
print(str(train_model.param_init_net.Proto())[:400] + '\n...')
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()))
def plot_mini_net(file_name, net_proto):
#plot the net
graph = net_drawer.GetPydotGraphMinimal(net_proto.Proto().op,
rankdir="LR",
minimal_dependency=True)
graph.write_png(file_name)
def TrainTest(args):
if not os.path.exists(args.save_folder):
os.makedirs(args.save_folder)
np.random.seed(123) # make test deterministic
arg_scope = {"order": "NCHW"}
train_model = model_helper.ModelHelper(name="mnist_train", arg_scope=arg_scope)
data, label = AddInput(
train_model,
batch_size=128,
db=os.path.join(args.data_folder, 'mnist-train-nchw-lmdb'),
db_type='lmdb')
loss, _ = AddLeNetModel(train_model, data, label, margin=args.margin)
AddTrainingOperators(train_model, loss)
# 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(args.data_folder, 'mnist-test-nchw-lmdb'),
db_type='lmdb')
AddLeNetModel(test_model, data, label, margin=args.margin)
# 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, no_loss=True, margin=args.margin)
# 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.
# 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
# print (str(train_model.param_init_net.Proto()))
graph = net_drawer.GetPydotGraphMinimal(train_model.net.Proto(),
"mnist",
rankdir="LR",
minimal_dependency=True)
graph.write(os.path.join(args.save_folder, "mnist.pdf"), format='pdf')
total_iters = 15000
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')
if i % 10 == 0:
print("#Iteration: {}, lambda: {:.3f}, loss: {:.3f}, train_acc: {:.3f}".
format(i, workspace.FetchBlob('lambda').tolist(), loss[i], accuracy[i]))
# 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)
embeds = []
labels = []
for i in range(100):
workspace.RunNet(test_model.net.Proto().name)
test_accuracy[i] = workspace.FetchBlob('accuracy')
embeds.append(workspace.FetchBlob('embedding'))
labels.append(workspace.FetchBlob('label'))
embeds = np.vstack(embeds)
labels = np.hstack(labels)
# vis, plot code from https://github.com/pangyupo/mxnet_center_loss
num = len(labels)
names = dict()
for i in range(10):
names[i]=str(i)
palette = np.array(sns.color_palette("hls", 10))
f = plt.figure(figsize=(8, 8))
ax = plt.subplot(aspect='equal')
sc = ax.scatter(embeds[:,0], embeds[:,1], lw=0, s=40,
c=palette[labels.astype(np.int)])
ax.axis('off')
ax.axis('tight')
# We add the labels for each digit.
txts = []
for i in range(10):
# Position of each label.
xtext, ytext = np.median(embeds[labels == i, :], axis=0)
txt = ax.text(xtext, ytext, names[i])
txt.set_path_effects([PathEffects.Stroke(linewidth=5, foreground="w"),
PathEffects.Normal()])
txts.append(txt)
fname = "distance-margin-{}.png".format(args.margin)
plt.savefig(os.path.join(args.save_folder, fname))
fig2 = plt.figure()
ax2 = fig2.add_subplot(111)
ax2.plot(loss, 'b')
ax2.plot(accuracy, 'r')
ax2.legend(('Loss', 'Accuracy'), loc='upper right')
plt.show()
fname = "loss-margin-{}.png".format(args.margin)
plt.savefig(os.path.join(args.save_folder, fname))
print('test_accuracy: %f' % test_accuracy.mean())
def GetPydotGraphMinimal(net, rankdir="BT"):
return net_drawer.GetPydotGraphMinimal(net, rankdir="BT")
def AddTrainingOperators(model, softmax, label):
# something very important happens here
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 - CNNModelHelper 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)
# let's checkpoint every 20 iterations, which should probably be fine.
# you may need to delete tutorial_files/tutorial-mnist to re-run the tutorial
model.Checkpoint([ITER] + model.params, [],
db="mnist_lenet_checkpoint_%05d.leveldb",
db_type="leveldb",
every=20)
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-leveldb'),
db_type='leveldb')
softmax = AddLeNetModel(train_model, data)
AddTrainingOperators(train_model, softmax, label)
# 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-leveldb'),
db_type='leveldb')
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")
graph = net_drawer.GetPydotGraphMinimal(train_model.net.Proto().op,
"mnist",
rankdir="LR",
minimal_dependency=True)
display.Image(graph.create_png(), width=800)
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)
# 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)
# 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.Proto().name)
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')
# 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')
# run a test pass on the test net
workspace.RunNetOnce(test_model.param_init_net)
workspace.CreateNet(test_model.net)
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())
