def load_nets(args, cur_gpu):
# initialize solver and feature net,
# RNN should be initialized before CNN, because CNN cudnn conv layers
# may assume using all available memory
caffe.set_mode_gpu()
caffe.set_device(cur_gpu)
solver = caffe.SGDSolver(args.solver)
if args.snapshot:
print "Restoring history from {}".format(args.snapshot)
solver.restore(args.snapshot)
net = solver.net
if args.weights:
print "Copying weights from {}".format(args.weights)
net.copy_from(args.weights)
return solver, net
评论列表
文章目录