resnet_imagenet_model_multi_wrapper.py 文件源码

def build(self, inp):
        # Divide input equally.
        self.lazy_init_var()
        inp_list = []
        output = []
        for ii in xrange(self.num_replica):
            with tf.name_scope('%s_%d' % ('replica', ii)) as scope:
                device = '/gpu:{}'.format(ii)
                with tf.device(device):
                    tf.get_variable_scope().reuse_variables()
                    inp_ = {
                        'x': inp['x_{}'.format(ii)],
                        'y_gt': inp['y_gt_{}'.format(ii)],
                        'phase_train': inp['phase_train']
                    }
                    output.append(self.sub_models[ii].build(inp_))
                    inp_list.append(inp_)
        self.output_list = output
        self.input_list = inp_list
        output = tf.concat(0, [oo['y_out'] for oo in output])
        self.register_var('y_out', output)
        output2 = tf.concat(0, [mm.get_var('score_out')
                                for mm in self.sub_models])
        self.register_var('score_out', output2)
        return {'y_out': output}