def _set_initializers()

in lm_human_preferences/rewards.py [0:0]

• 20 lines of code
• 8 McCabe index (conditional complexity)

    def _set_initializers(self):
        """Change initializers to load a model from a tensorflow checkpoint."""
        if self.comm.Get_rank() > 0 or self.train_dir == 'test':
            return

        assert self.model.built
        checkpoint_scope = 'reward_model'

        with tf.init_scope():
            # Initialize!
            params = {v.op.name: v for v in self.get_params()}
            checkpoint = tf.train.latest_checkpoint(os.path.join(self.train_dir, 'checkpoints/'))
            available = tf.train.list_variables(checkpoint)
            unchanged = {}

            for name, shape in available:
                if not name.startswith(checkpoint_scope + '/'):
                    # print('skipping', name)
                    continue
                if name.endswith('adam') or name.endswith('adam_1'):
                    # print('skipping', name)
                    continue
                print('setting', name)
                var = params[self.scope + name[len(checkpoint_scope):]]
                assert var.shape == shape, 'Shape mismatch: %s.shape = %s != %s' % (var.op.name, var.shape, shape)
                unchanged[name] = var
            tf.train.init_from_checkpoint(checkpoint, unchanged)