def _calculate_lambda_closeness()

in domainbed_measures/measures/held_out_measures.py [0:0]


    def _calculate_lambda_closeness(self,
                                    all_train_feats,
                                    all_train_labels,
                                    all_held_out_feats,
                                    all_held_out_labels,
                                    lr_decay_gamma,
                                    num_head_batches,
                                    max_lr,
                                    lr_sweep_factor,
                                    train_env_to_use,
                                    train_val_split,
                                    trainval_test_split=0.8):

        # Permute the datapoints from train feats
        all_train_feats, all_train_labels = permute_dataset(
            all_train_feats, all_train_labels)
        all_held_out_feats, all_held_out_labels = permute_dataset(
            all_held_out_feats, all_held_out_labels)

        ndata = min(all_train_feats.shape[0], all_held_out_feats.shape[0])
        all_train_feats = all_train_feats[:ndata, :]
        all_train_labels = all_train_labels[:ndata]
        all_held_out_feats = all_held_out_feats[:ndata, :]
        all_held_out_labels = all_held_out_labels[:ndata]

        all_feats = torch.vstack([all_train_feats, all_held_out_feats])
        all_labels = torch.hstack([all_train_labels, all_held_out_labels])
        all_feats, all_labels = permute_dataset(all_feats, all_labels)

        feat_dim = all_train_feats.shape[-1]

        logging.info("Obtaining heads")
        callbacks = [
            skorch.callbacks.LRScheduler(
                torch.optim.lr_scheduler.StepLR,
                gamma=lr_decay_gamma,
                step_size=self._train_epochs / 2,
            ),
            skorch.callbacks.EpochScoring(
                self.accuracy_fn,
                lower_is_better=False,
                name='val_accuracy',
            ),
            skorch.callbacks.EpochScoring(
                self.accuracy_fn,
                lower_is_better=False,
                name='train_accuracy',
                on_train=True,
            ),
            skorch.callbacks.EarlyStopping(
                monitor='val_accuracy',
                patience=15,
                threshold=0.0001,
                threshold_mode='rel',
                lower_is_better=False,
            )
        ]

        heads = self.get_heads(
            num_head_batches,
            feat_dim=feat_dim,
            criterion=nn.CrossEntropyLoss,
            num_labels=len(torch.unique(all_train_labels)),
            max_lr=max_lr,
            lr_sweep_factor=lr_sweep_factor,
            train_split=skorch.dataset.CVSplit(train_val_split),
            batch_size=self._algorithm.hparams['batch_size'],
            callbacks=callbacks)

        val_accuracies = []
        for hidx, h in enumerate(heads):
            logging.info("Fitting head %d/%d" % (hidx, len(heads)))
            train_val_idx = int(trainval_test_split * all_feats.shape[0])
            h.fit(all_feats[:train_val_idx, :], all_labels[:train_val_idx])
            val_accuracies.append([x['val_accuracy'] for x in h.history][-1])

        best_model_idx = np.argmax(val_accuracies)
        best_accuracy = heads[best_model_idx].accuracy(
            all_feats[train_val_idx:, :], all_labels[train_val_idx:])

        return 2.0 * (2.0 * best_accuracy - 1)