caffe 模型训练

import caffe
from caffe import layers as L
import os
import numpy as np
from matplotlib import pyplot as plt

Solver_PATH = 'auto_solver_.prototxt'

def change_env():
	root = os.path.dirname(__file__)
	os.chdir(root)
	print ("current work root:->",root)

def net(img_list,batch_size,mean_value=0):
	network = caffe.NetSpec()
	network.data,network.label = L.ImageData(source=img_list,batch_size=batch_size,new_width=28,new_height=28,ntop=2, transform_param=dict(scale=1/255.0, mean_value=mean_value))
	network.ip1  = L.InnerProduct(network.data,num_output=50,weight_filler=dict(type="xavier"))
	network.relu1 = L.ReLU(network.ip1,in_place=True)
	network.ip2  = L.InnerProduct(network.relu1,num_output=10,weight_filler=dict(type="xavier"))
	network.loss = L.SoftmaxWithLoss(network.ip2,network.label)
	network.accu = L.Accuracy(network.ip2,network.label,include={'phase':caffe.TEST})
	return network.to_proto()


def file_write(path1="auto_train.prototxt",path2="auto_test.prototxt"):
	with open(path1,"w") as f:
		f.write(str(net("train.imglist",200,108)))
	with open(path2,"w") as f:
		f.write(str(net("test.imglist",50,108)))

def main():
	#change_env()
	file_write()
	#caffe.set_device(0)
	#caffe.set_mode_gpu()
	solver = caffe.SGDSolver(Solver_PATH)
	solver.solve()
	
	iternum = 401

	loss_train = np.zeros(iternum)
	accu_train = np.zeros(iternum)
	accu_test  = np.zeros(iternum)

	for it in range(iternum):
	    solver.step(10) 
	    loss_train[it] = solver.net.blobs['loss'].data
	    accu_train[it] = solver.net.blobs['accu'].data
	    accu_test[it] = solver.test_nets[0].blobs['accu'].data
	    

	fig, ax = plt.subplots()
	
	ax.plot(np.arange(iternum),loss_train,c="r")
	ax.plot(np.arange(iternum),accu_train,c="g")
	ax.plot(np.arange(iternum),accu_test,c="b")

	ax.set_xlabel("iteration")
	ax.set_ylabel("Train loss")

	plt.show()
	


if __name__ == '__main__':
	main()

完整项目：https://download.csdn.net/download/qq_38641985/18434137