prune.py 文件源码

def retrain_pruned(solver, pruned_caffemodel, threshold, prune_layers):
    #solver = caffe.SGDSolver(solver_proto)
    retrain_iter = 20

    accuracys = []
    for i in range(retrain_iter):
        solver.net.copy_from(pruned_caffemodel)
        # solver.solve()
        solver.step(500)
        _,_,_,_,accuracy=prune(threshold, solver.test_nets[0], prune_layers)
        solver.test_nets[0].save(pruned_caffemodel)
        accuracys.append(accuracy)

    plt.plot(accuracys, 'r.-')
    plt.show()


#CPU?GPU????
#caffe.set_mode_cpu()