setup_predictor_towers.py 文件源码

def __init__(self, conf, gpu_id, start_images, actions, start_states, pix_distrib1,pix_distrib2):
        nsmp_per_gpu = conf['batch_size']/ conf['ngpu']

        # picking different subset of the actions for each gpu
        startidx = gpu_id * nsmp_per_gpu
        actions = tf.slice(actions, [startidx, 0, 0], [nsmp_per_gpu, -1, -1])
        start_images = tf.slice(start_images, [startidx, 0, 0, 0, 0], [nsmp_per_gpu, -1, -1, -1, -1])
        start_states = tf.slice(start_states, [startidx, 0, 0], [nsmp_per_gpu, -1, -1])

        pix_distrib1 = tf.slice(pix_distrib1, [startidx, 0, 0, 0, 0], [nsmp_per_gpu, -1, -1, -1, -1])
        pix_distrib2 = tf.slice(pix_distrib2, [startidx, 0, 0, 0, 0], [nsmp_per_gpu, -1, -1, -1, -1])

        print 'startindex for gpu {0}: {1}'.format(gpu_id, startidx)

        from prediction_train_sawyer import Model

        if 'ndesig' in conf:
            self.model = Model(conf, start_images, actions, start_states, pix_distrib=pix_distrib1,pix_distrib2=pix_distrib2, inference=True)
            # self.model = Model(conf, start_images, actions, start_states, pix_distrib=pix_distrib1,
            #                    pix_distrib2=pix_distrib2,
            #                    reuse_scope=reuse_scope)
        else:
            # self.model = Model(conf,start_images,actions,start_states, pix_distrib=pix_distrib1, reuse_scope= reuse_scope)
            self.model = Model(conf, start_images, actions, start_states, pix_distrib=pix_distrib1, inference=True)