python类set_device()的实例源码

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """

    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)
    # set up caffe
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)

def __init__(self, mean, weight, K, num_act, num_step=1, data_path='test'):
        self.K = K
        self.num_act = num_act
        self.num_step = num_step

        caffe.set_mode_gpu()
        caffe.set_device(0)

        test_net_file, net_proto = N.create_netfile(1, data_path, mean, K, K,
            1, num_act, num_step=self.num_step, mode='test')

        self.test_net = caffe.Net(test_net_file, caffe.TEST)
        self.test_net.copy_from(weight)

def __init__(self, mean, weight, K, num_act, num_step=1, data_path='test'):
        self.K = K
        self.num_act = num_act
        self.num_step = num_step

        caffe.set_mode_gpu()
        caffe.set_device(0)

        test_net_file, net_proto = N.create_netfile(1, data_path, mean, K, K,
            1, num_act, num_step=self.num_step, mode='test')

        self.test_net = caffe.Net(test_net_file, caffe.TEST)
        self.test_net.copy_from(weight)

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """

    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)
    # set up caffe
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)

def ctpnSource():
    DEMO_IMAGE_DIR = "img/"
    NET_DEF_FILE = "CTPN/models/deploy.prototxt"
    MODEL_FILE = "CTPN/models/ctpn_trained_model.caffemodel"
    caffe.set_mode_gpu()
    caffe.set_device(cfg.TEST_GPU_ID)
    # initialize the detectors
    text_proposals_detector = TextProposalDetector(CaffeModel(NET_DEF_FILE, MODEL_FILE))
    text_detector = TextDetector(text_proposals_detector)
    return text_detector

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('dataset', nargs='?',
                        choices=['pascal_voc', 'camvid', 'kitti', 'cityscapes'])
    parser.add_argument('input_path', nargs='?', default='',
                        help='Required path to input image')
    parser.add_argument('-o', '--output_path', default=None)
    parser.add_argument('--gpu', type=int, default=-1,
                        help='GPU ID to run CAFFE. '
                             'If -1 (default), CPU is used')
    args = parser.parse_args()
    if args.input_path == '':
        raise IOError('Error: No path to input image')
    if not exists(args.input_path):
        raise IOError("Error: Can't find input image " + args.input_path)
    if args.gpu >= 0:
        caffe.set_mode_gpu()
        caffe.set_device(args.gpu)
        print('Using GPU ', args.gpu)
    else:
        caffe.set_mode_cpu()
        print('Using CPU')
    if args.output_path is None:
        args.output_path = '{}_{}.png'.format(
                splitext(args.input_path)[0], args.dataset)
    predict(args.dataset, args.input_path, args.output_path)

def main(argv):
    model_filename = ''
    weight_filename = ''
    img_filename = ''
    try:
        opts, args = getopt.getopt(argv, "hm:w:i:")
        print opts
    except getopt.GetoptError:
        print 'yolo_main.py -m <model_file> -w <output_file> -i <img_file>'
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-h':
            print 'yolo_main.py -m <model_file> -w <weight_file> -i <img_file>'
            sys.exit()
        elif opt == "-m":
            model_filename = arg
        elif opt == "-w":
            weight_filename = arg
        elif opt == "-i":
            img_filename = arg
    print 'model file is "', model_filename
    print 'weight file is "', weight_filename
    print 'image file is "', img_filename

    caffe.set_device(0)
    caffe.set_mode_gpu()
    net = caffe.Net(model_filename, weight_filename, caffe.TEST)
    img = caffe.io.load_image(img_filename)  # load the image using caffe io
    img_ = scipy.misc.imresize(img, (448, 448))
    transformer = SimpleTransformer([104.00699, 116.66877, 122.67892])
    input = transformer.preprocess(img_)
    out = net.forward_all(data=input)
    print out.iteritems()
    img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    results = interpret_output(out['result'][0], img.shape[1], img.shape[0])  # fc27 instead of fc12 for yolo_small
    show_results(img_cv, results, img.shape[1], img.shape[0])
    cv2.waitKey(0)

def LoadNet(self,model,weights):
        caffe.set_mode_gpu()
        caffe.set_device(0)
        Net = caffe.Net(model, weights, caffe.TEST)
        return Net

def net():
    """Delay loading the net until the last possible moment.

    Loading the net is SLOW and produces a ton of terminal garbage.
    Also we want to wait to load it until we have called some other
    caffe initializations code (caffe.set_mode_gpu(), caffe.set_device(0), etc)

    """
    global __net
    if __net is None:
        __net = caffe.Net(LAYERS, WEIGHTS, caffe.TEST)
    return __net

def load_models(args):

    # load rnn model
    config = TestConfig()
    config.num_layers = args.lstm_num
    config.type = args.lstm_type
    config.hidden_size = config.input_size = args.lstm_input_size

    #tf.set_random_seed(1017)
    sess_config = tf.ConfigProto()
    # sess_config.gpu_options.allow_growth=True
    with tf.Graph().as_default():
        session = tf.Session(config=sess_config)
        initializer = tf.random_uniform_initializer(-config.init_scale,
                                                    config.init_scale)
        with tf.variable_scope("model", reuse=None, initializer=None):
            # with tf.device("/gpu:{}".format(args.job_id)):
            with tf.device("/cpu:0"):
                print "Constructing RNN network..."
                rnn_net = TPNModel(is_training=False, config=config)

        # restoring variables
        saver = tf.train.Saver()
        print "Starting loading session..."
        saver.restore(session, args.lstm_path)
        print 'Loaded RNN network from {:s}.'.format(args.lstm_path)

    # load feature model
    caffe.set_mode_gpu()
    caffe.set_device(args.job_id - 1)
    feature_net = caffe.Net(args.def_file, args.param, caffe.TEST)
    print 'Loaded feature network from {:s}.'.format(args.def_file)

    return feature_net, rnn_net, session

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)
    net = caffe.Net(args.model, args.weights, caffe.TEST)
    return net

def train_net(solver_prototxt, roidb, output_dir, nccl_uid, gpus, rank,
        queue, bbox_means, bbox_stds, pretrained_model=None, max_iters=40000):
    """Train a Fast R-CNN network."""
    caffe.set_mode_gpu()
    caffe.set_device(gpus[rank])
    caffe.set_solver_count(len(gpus))
    caffe.set_solver_rank(rank)
    caffe.set_multiprocess(True)
    caffe.set_random_seed(cfg.RNG_SEED)
    sw = SolverWrapper(solver_prototxt, roidb, output_dir, nccl_uid, 
        rank, bbox_means, bbox_stds, pretrained_model=pretrained_model)
    model_paths = sw.train_model(max_iters)
    if rank==0:
        queue.put(model_paths)

def rpn_generate_single_gpu(prototxt, caffemodel, imdb, rank, gpus, output_dir):
    cfg.GPU_ID = gpus[rank]
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)
    net = caffe.Net(prototxt, caffemodel, caffe.TEST)
    imdb_boxes = imdb_proposals(net, imdb, rank, len(gpus), output_dir)

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """

    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)
    # set up caffe
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)

def prep_for_deploy(self, batch_size, source_net=False, target_net=False, deploy_fn='deploy.proto', caffemodel_fn='score.caffemodel', gpu_id=0):
        caffe.set_mode_gpu()
        caffe.set_device(gpu_id)

        self.generate_deploy_proto(deploy_fn, batch_size, source_net=source_net, target_net=target_net)
        self.deploy = caffe.Net(deploy_fn, caffe.TEST, weights=caffemodel_fn)

        self._set_semantics(self.deploy, source=False, init_cw=False)
        self._set_semantics(self.deploy, source=True, init_cw=False)

def prep_for_training(self, model_fn, solver_fn, dt_lmdbs, sem_lmdbs, trainOpts, batch_size, gpu_id):
        caffe.set_mode_gpu()
        caffe.set_device(gpu_id)

        self.generate_train_proto(model_fn, dt_lmdbs, sem_lmdbs, batch_size)
        self.generate_solver_proto(solver_fn, model_fn, trainOpts=trainOpts)

        solver = caffe.NesterovSolver(solver_fn)
        self.base_cnn.load_pretrained(solver.net)
        self._set_semantics(solver.net, source=True, init_cw=True)
        self._set_semantics(solver.test_nets[1], source=False, init_cw=True)

        self.solver = solver

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """

    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)
    # set up caffe
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)

def __init__(self, path_to_deploy_file, path_to_model_file, input_layer_name="data_q", gpu_mode=True, device_id=1,
                 height=None, width=None):
        self.path_to_deploy_file = path_to_deploy_file
        self.path_to_model_file = path_to_model_file
        if gpu_mode:
            caffe.set_mode_gpu()
            caffe.set_device(device_id)
        else:
            caffe.set_mode_cpu()
        self.net = caffe.Net(path_to_deploy_file, path_to_model_file, caffe.TEST)
        self.input_layer_name = input_layer_name
        self.height = height or self.net.blobs[self.input_layer_name].data.shape[2]
        self.width = width or self.net.blobs[self.input_layer_name].data.shape[3]

def _init_caffe(cfg):
    """Initialize pycaffe in a training process.
    """

    import caffe
    # fix the random seeds (numpy and caffe) for reproducibility
    np.random.seed(cfg.RNG_SEED)
    caffe.set_random_seed(cfg.RNG_SEED)
    # set up caffe
    caffe.set_mode_gpu()
    caffe.set_device(cfg.GPU_ID)

def start(self, rank):
        self.rank = rank

        if len(self.gpus) > 0:
            self.device = self.gpus[rank]
            if debug:
                s = 'solver gpu %d' % self.gpus[self.rank] + \
                    ' pid %d' % os.getpid() + ' size %d' % self.size + \
                    ' rank %d' % self.rank
                print(s, file = sys.stderr)
            caffe.set_mode_gpu()
            caffe.set_device(self.device)
            caffe.set_solver_count(self.size)
            caffe.set_solver_rank(self.rank)
            caffe.set_multiprocess(True)
        else:
            print('solver cpu', file = sys.stderr)
            caffe.set_mode_cpu()

        if self.cmd.graph.endswith('.json'):
            with open(self.cmd.graph, mode = 'r') as f:
                graph = caffe_pb2.SolverParameter()
                text_format.Merge(f.read(), graph)
                self.graph = graph
        else:
            self.graph = self.solver_graph()

        import tempfile
        with tempfile.NamedTemporaryFile(mode = 'w+', delete = False) as f:
            text_format.PrintMessage(self.graph, f)
            tmp = f.name
        self.caffe = caffe.AdamSolver(tmp)

        if self.uid:
            self.nccl = caffe.NCCL(self.caffe, self.uid)
            self.nccl.bcast()
            self.caffe.add_callback(self.nccl)
            if self.caffe.param.layer_wise_reduce:
                self.caffe.net.after_backward(self.nccl)