# 训练设置 # 使用GPU caffe.set_device(gpu_id) # 若不设置,默认为0 caffe.set_mode_gpu() # 使用CPU caffe.set_mode_cpu() # 加载Solver，有两种常用方法 # 1. 无论模型中Slover类型是什么统一设置为SGD solver = caffe.SGDSolver('/home/xxx/data/solver.prototxt') # 2. 根据solver的prototxt中solver_type读取，默认为SGD solver = caffe.get_solver('/home/xxx/data/solver.prototxt') # 训练模型 # 1.1 前向传播 solver.net.forward() # train net solver.test_nets[0].forward() # test net (there can be more than one) # 1.2 反向传播,计算梯度 solver.net.backward() # 2. 进行一次前向传播一次反向传播并根据梯度更新参数 solver.step(1) # 3. 根据solver文件中设置进行完整model训练 solver.solve()
如果想在训练过程中保存模型参数，调用
solver.net.save('mymodel.caffemodel')