train_MFISNet_Refinement.py

"""
This script does the following:
1. Loads data from specified directories
2. Sets up an ParallelFYNetInverse model (with or without a skip connection)
3. Trains the model
4. Runs inference on all of the train / val samples.
5. Transforms the intermediate outputs to Cart coordinates.
6. Writes the scattering objects to disk
"""

import logging
from typing import Tuple, Dict, List, Callable
import argparse
from timeit import default_timer
import os
import sys
import shlex
import numpy as np
import h5py
import torch
from timeit import default_timer
import pandas as pd
import wandb

from src.data.data_naming_constants import (
    X_VALS,
    RHO_VALS,
    THETA_VALS,
    Q_POLAR,
    Q_POLAR_LPF,
    Q_CART,
    D_MH,
    KEYS_FOR_TRAINING_METADATA,
    D_RS,
    SAMPLE_COMPLETION,
    H_VALS,
    M_VALS,
    KEYS_FOR_EXPERIMENT_INFO_OUT,
)


from src.data.data_io import load_dir
from src.models.MFISNet_Refinement import (
    MFISNet_Refinement,
)
from src.training_utils.train_loop import (
    train,
    evaluate_losses_on_dataloader,
    EarlyStopper,
)
from src.training_utils.loss_functions import (
    MSEModule,
)
from src.training_utils.make_predictions import make_preds_on_dataset
from src.utils.logging_utils import FMT, TIMEFMT, write_result_to_file, hash_dict
from src.data.add_noise import add_noise_to_d


def setup_args(argument_string: str = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser()

    parser.add_argument("-meas_dir_train_frmt")
    parser.add_argument("-scobj_dir_train")
    parser.add_argument("-meas_dir_val_frmt")
    parser.add_argument("-scobj_dir_val")
    parser.add_argument("-output_dir_train")
    parser.add_argument("-output_dir_val")
    parser.add_argument("-results_fp")
    parser.add_argument("-model_weights_dir")
    parser.add_argument("-wavenumbers", nargs="+")
    parser.add_argument("-truncate_num", type=int)
    parser.add_argument("-truncate_num_val", type=int)
    parser.add_argument("-seed", type=int, default=35675)
    parser.add_argument("-n_cnn_1d", type=int, default=3)
    parser.add_argument("-n_cnn_2d", type=int, default=3)
    parser.add_argument("-n_cnn_channels_1d", type=int, default=10)
    parser.add_argument("-n_cnn_channels_2d", type=int, default=10)
    parser.add_argument("-kernel_size_1d", type=int, default=13)
    parser.add_argument("-kernel_size_2d", type=int, default=13)
    parser.add_argument("-batch_size", type=int, default=32)
    parser.add_argument("-n_epochs_pretrain", type=int, default=10)
    parser.add_argument("-n_epochs_finetune", type=int, default=10)
    parser.add_argument("-lr_init", type=float, default=1.0)
    parser.add_argument("-weight_decay", type=float, default=0.0)
    parser.add_argument("-eta_min", type=float, default=1e-04)
    parser.add_argument("-n_epochs_per_log", type=int, default=5)
    parser.add_argument("-debug", default=False, action="store_true")
    parser.add_argument("-wandb_project_name", type=str, help="W&B project name")
    parser.add_argument("-wandb_entity", type=str, help="The W&B entity")
    parser.add_argument("-wandb_mode", choices=["offline", "online"], default="offline")
    parser.add_argument("-dont_use_wandb", default=False, action="store_true")
    parser.add_argument("-skip_reader_setup", default=False, action="store_true")
    parser.add_argument("-dont_write_outputs", default=False, action="store_true")
    parser.add_argument("-input_q_hat", default=False, action="store_true")
    parser.add_argument("-input_after_adj", default=False, action="store_true")
    parser.add_argument("-warm_start_init", default=False, action="store_true")
    parser.add_argument("-lr_decrease_factor", default=None, type=float)
    parser.add_argument("-noise_to_signal_ratio", default=None, type=float)
    parser.add_argument(
        "-init_mode",
        default="original",
        choices=[
            "original",
            "uniform-with-old-scale",
            "normal-with-old-scale",
            "he-normal",
        ],
    )
    parser.add_argument("-train_without_filters", default=False, action="store_true")
    parser.add_argument("-early_stopping", default=False, action="store_true")
    parser.add_argument("-min_delta", type=float, default=1e-03)

    if argument_string is None:
        # Parse the arguments from the system's argv
        args = parser.parse_args()
    else:
        # Split the argument string into a list using shlex.split()
        arg_list = shlex.split(argument_string)
        args = parser.parse_args(arg_list)
        # Adding a hash because this doesn't exist in the argument string but is
        # necessary in the main function
        id_hash = hash_dict(vars(args))
        args.hash = id_hash

    return args


class MultiFreqResidData(torch.utils.data.Dataset):
    def __init__(
        self,
        d_mh: torch.Tensor,
        q_polar_lpf: torch.Tensor,
        output_nan_samples_bool: bool = False,
        freq_output_idx: int = None,
    ) -> None:
        """_summary_

        Args:
            d_mh (torch.Tensor): Has shape (N_samples, N_freqs, N_m, N_h)
            q_polar_lpf (torch.Tensor): Has shape (N_samples, N_freqs, N_m, N_h)
            output_nan_samples_bool (bool, optional): Defaults to False.
        """
        self.d_mh = torch.view_as_real(d_mh).to(torch.float)
        self.q_polar_lpf = q_polar_lpf.to(torch.float)
        self.output_nan_samples_bool = output_nan_samples_bool
        logging.debug(
            "MultiFreqResidData: q_polar_lpf shape: %s", self.q_polar_lpf.shape
        )
        logging.debug("MultiFreqResidData: d_mh shape: %s", self.d_mh.shape)

        nan_samples = torch.any(torch.isnan(self.d_mh[:, :, 0, 0, 0]), dim=1)
        self.keep_indices = torch.argwhere(torch.logical_not(nan_samples)).flatten()
        logging.debug(
            "MultiFreqResidData: nan_samples has shape: %s  and sum %s",
            nan_samples.shape,
            torch.sum(nan_samples),
        )
        logging.debug(
            "MultiFreqResidData: self.keep_indices: shape: %s", self.keep_indices.shape
        )

        logging.info(
            "Initialized a MultiFreqResidData instance with d_mh shape: %s and q_polar_lpf shape: %s",
            self.d_mh.shape,
            self.q_polar_lpf.shape,
        )

        self.n_samples = self.d_mh.shape[0]
        if freq_output_idx is None:
            self.freq_output_idx = 0
        else:
            self.freq_output_idx = freq_output_idx

    def __len__(self):
        if self.output_nan_samples_bool:
            return self.n_samples
        else:
            return self.keep_indices.shape[0]

    def __getitem__(self, idx):
        """Returns (x, y, z)
        where:
         - x is the wave field d_mh at multiple different freqs. It has shape (N_M, N_H, N_freqs)
         - y is the target scattering object in polar coords. It
        """
        if not self.output_nan_samples_bool:
            idx = self.keep_indices[idx]

        x = self.d_mh[idx]
        return x, self.q_polar_lpf[idx, self.freq_output_idx], self.q_polar_lpf[idx, -1]

    def __repr__(self) -> str:
        _, N_freqs, N_m, N_h, _ = self.d_mh.shape
        N_samples = self.__len__()
        s = f"MultiFreqResidData object with N_m={N_m}, N_h={N_h}, N_freqs={N_freqs}, freq_output_idx={self.freq_output_idx}"
        s += f"  N_samples={N_samples}"
        return s


def load_data(
    meas_data_dir_frmt: str,
    scobj_data_dir: str,
    wavenumbers: List[str],
    truncate_num: int = None,
    noise_to_sig_ratio: float = None,
    train_without_filters: bool = False,
) -> MultiFreqResidData:
    """This function loads all of the samples indicated by the data_dir.
    It arranges the data and returns a MultiFreqResidData object.

    Args:
        data_dir (str): Where the data is located
        truncate_num (int, optional): _description_. Defaults to None.

    Returns:
        MultiFreqResidData: _description_
    """

    dd_last = load_dir(
        meas_data_dir_frmt.format(wavenumbers[-1]),
        scobj_data_dir,
        truncate_num=truncate_num,
    )

    d_mh = np.empty(
        (
            dd_last[D_MH].shape[0],
            len(wavenumbers),
            dd_last[D_MH].shape[1],
            dd_last[D_MH].shape[2],
        ),
        dtype=np.complex64,
    )

    d_mh[:, -1] = dd_last[D_MH]

    q_polar = dd_last[Q_POLAR]
    q_polar_lpf = np.empty(
        (q_polar.shape[0], len(wavenumbers), q_polar.shape[1], q_polar.shape[2])
    )
    q_polar_lpf[:, -1] = q_polar

    # logging.debug(
    #     "load_data: Now I am looping over the remaining wavenumbers: %s",
    #     wavenumbers[:-1],
    # )

    for i, w in enumerate(wavenumbers[:-1]):
        # logging.debug("load_data: w=%s, i=%s", w, i)
        dd_i = load_dir(
            meas_data_dir_frmt.format(w),
            scobj_data_dir,
            truncate_num=truncate_num,
        )
        d_mh[:, i] = dd_i[D_MH]
        if train_without_filters:
            q_polar_lpf[:, i] = dd_i[Q_POLAR]
        else:
            q_polar_lpf[:, i] = dd_i[Q_POLAR_LPF]

    if noise_to_sig_ratio is not None:
        logging.info("Adding noise with noise-to-signal ratio %f", noise_to_sig_ratio)
        d_mh = add_noise_to_d(d_mh, noise_to_sig_ratio)

    q_polar_lpf = torch.from_numpy(q_polar_lpf)
    d_mh = torch.from_numpy(d_mh)

    dset_obj = MultiFreqResidData(d_mh, q_polar_lpf, output_nan_samples_bool=False)
    metadata_dd = {i: dd_last[i] for i in KEYS_FOR_EXPERIMENT_INFO_OUT}
    return dset_obj, metadata_dd


def main(args: argparse.Namespace) -> None:
    """
    1. Load data
    2. Do necessary transformations
    3. Set up NN
    4. Train NN
    """
    # Check arguments. Some settings of the arguments may not be acceptable. In this case, we'll
    # exit early.
    if args.input_after_adj and not args.input_q_hat:
        logging.warning(
            f"Invalid argument setting: input_q_hat={args.input_q_hat} and input_after_adj={args.input_after_adj}"
        )
        return

    if args.input_q_hat and args.warm_start_init:
        logging.warning(
            f"Invalid argument setting: input_q_hat={args.input_q_hat} and warm_star_init={args.warm_start_init}"
        )
        return

    # Make sure output directories exist
    d = [
        args.model_weights_dir,
    ]
    if args.output_dir_train is not None:
        d.append(args.output_dir_train)
    if args.output_dir_val is not None:
        d.append(args.output_dir_val)
    for i in d:
        if not os.path.isdir(i):
            os.mkdir(i)

    # Check whether this training run pre-exists
    if os.path.isfile(args.results_fp):
        pre_existing_bool = True
        results_df = pd.read_table(args.results_fp)
        last_row = results_df.iloc[-1]

        # Check to make sure all of the arguments are the same.
        assert (
            last_row["hash"] == args.hash
        ), f"Hashes don't match: {last_row['hash']} vs {args.hash}"

        final_epoch = last_row["epoch"]
        final_training_part = last_row["training_part"]
        if final_training_part.startswith("pretrain"):
            final_training_part_int = int(final_training_part.split("_")[-1])
        else:
            final_training_part_int = len(args.wavenumbers)
        logging.info(
            "Resuming after %i epochs during training part %s",
            final_epoch,
            final_training_part,
        )

    else:
        pre_existing_bool = False
        final_epoch = 0
        final_training_part = "pretrain_0"
        final_training_part_int = 0

    # Set seeds for reproducibility
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Load data from disk
    dset_train, train_metadata_dd = load_data(
        args.meas_dir_train_frmt,
        args.scobj_dir_train,
        args.wavenumbers,
        args.truncate_num,
        noise_to_sig_ratio=args.noise_to_signal_ratio,
        train_without_filters=args.train_without_filters,
    )
    dset_val, val_metadata_dd = load_data(
        args.meas_dir_val_frmt,
        args.scobj_dir_val,
        args.wavenumbers,
        args.truncate_num_val,
        noise_to_sig_ratio=args.noise_to_signal_ratio,
        train_without_filters=args.train_without_filters,
    )

    # Grab some metadata
    x_vals = train_metadata_dd[X_VALS]
    rho_vals = train_metadata_dd[RHO_VALS]
    theta_vals = train_metadata_dd[THETA_VALS]
    N_x = x_vals.shape[0]
    N_rho = rho_vals.shape[0]
    N_theta = theta_vals.shape[0]
    N_freqs = len(args.wavenumbers)
    logging.info("Working with %i wavenumbers: %s", N_freqs, args.wavenumbers)

    n_train = len(dset_train)
    n_val = len(dset_val)

    args.n_train = n_train
    args.n_val = n_val

    logging.info("Done Data Loading. N_train=%i, N_val=%i", n_train, n_val)

    #########################################################
    # Set dataloaders
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logging.info("Training on device: %s", device)

    #########################################################
    # Set up NN architecture

    model = MFISNet_Refinement(
        N_h=N_rho,
        N_rho=N_rho,
        c_1d=args.n_cnn_channels_1d,
        c_2d=args.n_cnn_channels_2d,
        w_1d=args.kernel_size_1d,
        w_2d=args.kernel_size_2d,
        N_cnn_1d=args.n_cnn_1d,
        N_cnn_2d=args.n_cnn_2d,
        N_freqs=len(args.wavenumbers),
        init_mode=args.init_mode,
    )

    if pre_existing_bool:
        fp_weights = os.path.join(args.model_weights_dir, f"epoch_{final_epoch}.pickle")
        logging.info("Loading weights from file: %s", fp_weights)
        incompatible_keys = model.load_state_dict(
            torch.load(fp_weights, map_location="cpu")
        )
        logging.info("Message from model loading: %s", incompatible_keys)

    ###########################################################################
    # SET UP LOG FUNCTION
    def log_function(model_0, epoch_local) -> bool:
        """
        NEED TO SET ANEW:
        loss_fn_dd
        epoch_stagger
        training_part
        n_epochs_this_step
        early_stopper (Optional)

        OTHER OUTER SCOPE STUFF:
        train_dloader,
        val_dloader,
        device,
        args
        """
        with torch.no_grad():
            epoch_eff = epoch_stagger + epoch_local

            # 1. Evaluate on train set
            train_loss_dd = evaluate_losses_on_dataloader(
                model_0, train_dloader, loss_fn_dd, device
            )
            # 2. Evaluate on val set
            val_loss_dd = evaluate_losses_on_dataloader(
                model_0, val_dloader, loss_fn_dd, device
            )

            weight_norm = torch.norm(
                torch.cat([x.view(-1) for x in model_0.parameters()]), 2
            )

            # 3. Log to console and log file
            logging.info(
                "%s Epoch %i/%i. Train MSE: %f, Train Rel L2: %f, Train PSNR: %f",
                training_part,
                epoch_local,
                n_epochs_this_step,
                torch.mean(train_loss_dd["mse"]).item(),
                torch.mean(train_loss_dd["rel_l2"]).item(),
                torch.mean(train_loss_dd["psnr"]).item(),
                # train_rel_l2_aaa,
                # train_psnr_aaa,
            )
            logging.info(
                "\t Val MSE: %f, Val Rel L2: %f, Val PSNR: %f",
                torch.mean(val_loss_dd["mse"]).item(),
                torch.mean(val_loss_dd["rel_l2"]).item(),
                torch.mean(val_loss_dd["psnr"]).item(),
            )
            logging.info("\t Weight L2 norm: %f", weight_norm.item())
            train_dd = {
                # Optimization info
                "epoch": epoch_local + epoch_stagger,
                "weight_norm": weight_norm.item(),
                "training_part": training_part,
                # Experiment info
                "n_train": n_train,
                "n_val": n_val,
                "n_cnn_1d": args.n_cnn_1d,
                "n_cnn_2d": args.n_cnn_2d,
                "n_cnn_channels_1d": args.n_cnn_channels_1d,
                "n_cnn_channels_2d": args.n_cnn_channels_2d,
                "kernel_size_1d": args.kernel_size_1d,
                "kernel_size_2d": args.kernel_size_2d,
                "lr_init": args.lr_init,
                "weight_decay": args.weight_decay,
                "batch_size": args.batch_size,
                "eta_min": args.eta_min,
                "n_rho_vals": N_rho,
                "n_theta_vals": N_theta,
                "hash": args.hash,
                "meas_dir_train_frmt": args.meas_dir_train_frmt,
                "scobj_dir_train": args.scobj_dir_train,
                "meas_dir_val_frmt": args.meas_dir_val_frmt,
                "scobj_dir_val": args.scobj_dir_val,
                "output_dir_train": args.output_dir_train,
                "output_dir_val": args.output_dir_val,
                "lr_decrease_factor": args.lr_decrease_factor,
                "init_mode": args.init_mode,
                "warm_start_init": args.warm_start_init,
            }
            for k, v in train_loss_dd.items():
                train_dd["train_" + k] = torch.mean(v).item()
            for k, v in val_loss_dd.items():
                train_dd["val_" + k] = torch.mean(v).item()
            write_result_to_file(args.results_fp, **train_dd)

            # Try to log results to W&B
            if not args.dont_use_wandb:
                try:
                    wandbrun.log(train_dd)
                except ValueError:
                    logging.error("Error: wandb logging failed for %s" % wandbrun.id)

        fp_weights = os.path.join(args.model_weights_dir, f"epoch_{epoch_eff}.pickle")
        torch.save(model_0.state_dict(), fp_weights)
        model_0 = model_0.to(device)

        # Check early stopping criterion
        if args.early_stopping:
            # Early stopping is based on the validation relative L2 error
            return early_stopper.early_stop(train_dd["val_rel_l2"])
        else:
            return False

    ##########################################################################################
    # PRETRAIN EACH BLOCK
    loss_module = MSEModule()
    loss_fn_dd = {
        "mse": loss_module.mse,
        "psnr": loss_module.psnr,
        "rel_l2": loss_module.relative_l2_error,
        "final_mse": loss_module.mse_against_final,
        "final_psnr": loss_module.psnr_against_final,
        "final_rel_l2": loss_module.relative_l2_error_against_final,
    }
    train_dloader = torch.utils.data.DataLoader(dset_train, batch_size=args.batch_size)
    val_dloader = torch.utils.data.DataLoader(dset_val, batch_size=args.batch_size)

    # Begin training timer
    t0 = default_timer()

    if args.n_epochs_pretrain:
        for i in range(final_training_part_int, N_freqs):
            training_part = f"pretrain_{i}"
            dset_train.freq_output_idx = i
            dset_val.freq_output_idx = i
            model.freq_pred_idx = i

            epoch_stagger = args.n_epochs_pretrain * i
            n_epochs_this_step = args.n_epochs_pretrain

            # Need to do some accounting to get the number of epochs right during the
            # pretraining part that was interrupted.
            if i == final_training_part_int and pre_existing_bool:
                epoch_stagger_old = epoch_stagger
                epoch_stagger = final_epoch + 1
                n_epochs_this_step = n_epochs_this_step - (
                    epoch_stagger - epoch_stagger_old
                )

            # SET ONLY THE Ith FYNET BLOCK TO BE TRAINABLE
            for j, block_j in enumerate(model.inverse_networks):
                if j == i:
                    for p in block_j.parameters():
                        p.requires_grad = True
                else:
                    for p in block_j.parameters():
                        p.requires_grad = False

            # Decrease the learning rate if desired
            if args.lr_decrease_factor is not None:
                lr_i = (args.lr_decrease_factor**i) * args.lr_init
                eta_min_i = (args.lr_decrease_factor**i) * args.eta_min
            else:
                lr_i = args.lr_init
                eta_min_i = args.eta_min

            # Optionally set up early stopping
            if args.early_stopping:
                early_stopper = EarlyStopper(min_delta=args.min_delta)
            model = train(
                model=model,
                n_epochs=n_epochs_this_step,
                lr_init=lr_i,
                weight_decay=args.weight_decay,
                momentum=0.0,
                eta_min=eta_min_i,
                train_loader=train_dloader,
                device=device,
                n_epochs_per_log=args.n_epochs_per_log,
                log_function=log_function,
                loss_function=loss_module,
            )

            # Now we want to find the best weights of this part of training and load these weights
            # back into the model.
            df = pd.read_table(args.results_fp)
            df_this_part = df[df["training_part"] == training_part]

            df_best = df_this_part[
                df_this_part["val_rel_l2"] == df_this_part["val_rel_l2"].min()
            ]

            best_epoch = df_best.iloc[0]["epoch"]

            logging.info("Epoch %i was best; loading these weights", best_epoch)

            best_weight_fp_i = os.path.join(
                args.model_weights_dir, f"epoch_{best_epoch}.pickle"
            )
            state_dict = torch.load(best_weight_fp_i, map_location=torch.device("cpu"))
            logging.debug("Now have state_dict with keys: %s", state_dict.keys())

            # Loop over the previous blocks, setting their weights
            for j in range(i + 1):
                key_str_j = f"inverse_networks.{j}."
                ll = len(key_str_j)
                state_dict_j = {
                    k[ll:]: v for k, v in state_dict.items() if k.startswith(key_str_j)
                }
                logging.debug(
                    "Now have state_dict_j with keys: %s", state_dict_j.keys()
                )
                model.inverse_networks[j].load_state_dict(state_dict_j)

                if j != 0:
                    # Also set the filtering blocks
                    key_str_j_filtering = f"filtering_blocks.{j-1}."
                    ll = len(key_str_j_filtering)
                    state_dict_j_filtering = {
                        k[ll:]: v
                        for k, v in state_dict.items()
                        if k.startswith(key_str_j_filtering)
                    }

                    model.filtering_blocks[j - 1].load_state_dict(
                        state_dict_j_filtering
                    )

            # If we're using the warm start weight initialization, use the weights of the current
            # block to initialize the next block.
            if args.warm_start_init and i + 1 < N_freqs:
                key_str_i = f"inverse_networks.{i}."
                ll = len(key_str_i)
                state_dict_i = {
                    k[ll:]: v for k, v in state_dict.items() if k.startswith(key_str_i)
                }
                logging.debug(
                    "Now have state_dict_i with keys: %s", state_dict_i.keys()
                )

                model.inverse_networks[i + 1].load_state_dict(state_dict_i)

            model = model.to(device)

    ##########################################################################################
    # FINE-TUNE THE WHOLE THING
    training_part = "fine_tune"
    loss_fn_dd = {
        "mse": loss_module.mse,
        "psnr": loss_module.psnr,
        "rel_l2": loss_module.relative_l2_error,
        "final_mse": loss_module.mse_against_final,
        "final_psnr": loss_module.psnr_against_final,
        "final_rel_l2": loss_module.relative_l2_error_against_final,
    }
    # Set all of the blocks trainable
    for block_j in model.inverse_networks:
        for p in block_j.parameters():
            p.requires_grad = True

    model.fine_tune_bool = True

    model.freq_pred_idx = N_freqs - 1
    dset_train.freq_output_idx = N_freqs - 1
    dset_val.freq_output_idx = N_freqs - 1

    epoch_stagger = args.n_epochs_pretrain * N_freqs
    n_epochs_this_step = args.n_epochs_finetune

    # Need to do some accounting to get the number of epochs right during the
    # fine-tuning part that was interrupted.
    if training_part == final_training_part and pre_existing_bool:
        epoch_stagger_old = epoch_stagger
        epoch_stagger = final_epoch + 1
        n_epochs_this_step = n_epochs_this_step - (epoch_stagger - epoch_stagger_old)

    if args.lr_decrease_factor is not None:
        lr_i = (args.lr_decrease_factor**N_freqs) * args.lr_init
        eta_min_i = (args.lr_decrease_factor**N_freqs) * args.eta_min
    else:
        lr_i = args.lr_init
        eta_min_i = args.eta_min

    # Optionally set up early stopping
    if args.early_stopping:
        early_stopper = EarlyStopper(min_delta=args.min_delta)
    model = train(
        model=model,
        n_epochs=n_epochs_this_step,
        lr_init=lr_i,
        weight_decay=args.weight_decay,
        momentum=0.0,
        eta_min=eta_min_i,
        train_loader=train_dloader,
        device=device,
        n_epochs_per_log=args.n_epochs_per_log,
        log_function=log_function,
        loss_function=loss_module,
    )

    t1 = default_timer()

    logging.info("Complete model training time is %f sec", t1 - t0)

    if not args.dont_write_outputs:
        #######################################################################
        # FIND EPOCH WITH MINIMAL TEST ERROR

        df = pd.read_table(args.results_fp)
        df_this_part = df[df["training_part"] == training_part]

        df_best = df_this_part[
            df_this_part["val_rel_l2"] == df_this_part["val_rel_l2"].min()
        ]

        best_epoch = df_best.iloc[0]["epoch"]

        logging.info("Epoch %i was best; loading these weights", best_epoch)

        best_weight_fp_final = os.path.join(
            args.model_weights_dir, f"epoch_{best_epoch}.pickle"
        )
        model.load_state_dict(
            torch.load(best_weight_fp_final, map_location=torch.device("cpu"))
        )
        model = model.to(device)

        #######################################################################
        # RUN INFERENCE ON ALL OF THE TRAIN/TEST SAMPLES

        dset_train.output_nans_bool = True
        dset_val.output_nans_bool = True

        n_train_all = len(dset_train)
        n_val_all = len(dset_val)

        train_dloader = torch.utils.data.DataLoader(
            dset_train, batch_size=args.batch_size
        )
        val_dloader = torch.utils.data.DataLoader(dset_val, batch_size=args.batch_size)
        logging.info(
            "Running inference on all %i train and %i val samples",
            n_train_all,
            n_val_all,
        )

        make_preds_on_dataset(
            model,
            train_dloader,
            device,
            args.output_dir_train,
            shard_size=500,
            experiment_info=train_metadata_dd,
        )

        make_preds_on_dataset(
            model,
            val_dloader,
            device,
            args.output_dir_val,
            shard_size=500,
            experiment_info=val_metadata_dd,
        )

    logging.info("Finished")


if __name__ == "__main__":
    a = setup_args()

    for name, logger in logging.root.manager.loggerDict.items():
        logging.getLogger(name).setLevel(logging.WARNING)

    if a.debug:
        logging.basicConfig(format=FMT, datefmt=TIMEFMT, level=logging.DEBUG)
    else:
        logging.basicConfig(format=FMT, datefmt=TIMEFMT, level=logging.INFO)

    id_hash = hash_dict(vars(a))
    a.hash = id_hash
    if a.dont_use_wandb:
        main(a)
    else:
        # print(vars(a))
        # print(id_hash)
        with wandb.init(
            id=id_hash,
            project=a.wandb_project_name,
            entity=a.wandb_entity,
            config=vars(a),
            mode=a.wandb_mode,
            reinit=True,
            resume=None,
            settings=wandb.Settings(start_method="fork"),
        ) as wandbrun:
            main(a)