您的位置:首页 > 编程语言 > Python开发

Creating caffe net prototxt files in python!

2016-04-05 23:55 435 查看
我们可以利用python来生成所有net结构的prototxt文件(当然,也可以直接写prototxt文件),这样,可以充分保证 training, testing和deploy 的一致性,并且,比较直观,修改也很方便

下面给出一个小例子

import caffe
from caffe import layers as L #pseudo module using __getattr__ magic to generate protobuf messages
from caffe import params as P #pseudo module using __getattr__ magic to generate protobuf messages

def setLayers(leveldb, batch_size, deploy=False):
#it is tricky to produce the deploy prototxt file, as the data input is not from a layer, so we have to creat a workaround
#producing training and testing prototxt files is pretty straight forward
n = caffe.NetSpec()
if deploy==False:
n.data, n.label = L.Data(batch_size=batch_size, backend=P.Data.LEVELDB, source=leveldb,\
ntop=2)
# produce data definition for deploy net
else:
input="data"
dim1=1
dim2=1
dim3=64
dim4=32
#make an empty "data" layer so the next layer accepting input will be able to take the correct blob name "data",
#we will later have to remove this layer from the serialization string, since this is just a placeholder
n.data=L.Layer()

n.conv2 = L.Convolution(n.data, kernel_h=6, kernel_w=1, num_output=8, weight_filler=dict(type='xavier'))
n.pool1 = L.Pooling(n.conv2, kernel_h=3, kernel_w=1, stride_h=2, stride_w=1, pool=P.Pooling.MAX)

n.drop2=L.Dropout(n.pool1,dropout_ratio=0.1)
n.ip1=L.InnerProduct(n.drop2, num_output=196, weight_filler=dict(type='xavier'))

n.relu1 = L.ReLU(n.ip1, in_place=True)
n.ip4 = L.InnerProduct(n.relu1, num_output=12, weight_filler=dict(type='xavier'))

#n.loss layer is only in training and testing nets, but not in deploy net.
if deploy==False:
n.loss = L.SoftmaxWithLoss(n.ip4, n.label)
return str(n.to_proto())
#for generating the deploy net
else:
#generate the input information header string
deploy_str='input: {}\ninput_dim: {}\ninput_dim: {}\ninput_dim: {}\ninput_dim: {}'.format('"'+input+'"', dim1, dim2, dim3, dim4)
#assemble the input header with the net layers string.  remove the first placeholder layer from the net string.
return deploy_str+'\n'+'layer {'+'layer {'.join(str(n.to_proto()).split('layer {')[2:])

#write the net prototxt files out
with open('trainNet.prototxt', 'w') as f:
print 'wrting train'
f.write(setLayers('databases/train_subj1_leveldb', 100))

with open('testNet.prototxt', 'w') as f:
print 'wrting test'
f.write(setLayers('databases/test_subj1_leveldb', 100))

with open('deploy.prototxt', 'w') as f:
f.write(str(setLayers('', 0, deploy=True)))
内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: