Ability to parallelize between GPUs

[matsen]

• factoring apart load_and_add_shm...
• finding least used gpu and using it, with a default
• branch lengths are tensors
• simplifying device handling of crepes

[matsen]

Dropped:

    def model_and_optimizer_to(self, device):
        self.model.to(device)
        for state in self.optimizer.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    state[k] = v.to(device)
        for dataset in [self.train_dataset, self.val_dataset]:
            if dataset is not None:
                dataset.to(device)

[matsen]

Also deferred to a future issue:

def dataset_of_pcp_df(pcp_df, branch_length_multiplier=5.0):
    return DNSMDataset.from_data(
        pcp_df["parents"],
        pcp_df["children"],
        pcp_df["rates"],
        pcp_df["subs_probs"],
        branch_length_multiplier=branch_length_multiplier,
    )


def train_val_datasets_of_pcp_df(pcp_df, branch_length_multiplier=5.0):
    """
    Perform a train-val split based on a "in_train" column.
    """
    train_df = pcp_df[pcp_df["in_train"]].reset_index(drop=True)
    val_df = pcp_df[~pcp_df["in_train"]].reset_index(drop=True)

    val_dataset = dataset_of_pcp_df(
        val_df,
        branch_length_multiplier=branch_length_multiplier,
    )
    if len(train_df) == 0:
        return None, val_dataset
    # else:
    train_dataset = dataset_of_pcp_df(
        train_df,
        branch_length_multiplier=branch_length_multiplier,
    )
    return train_dataset, val_dataset