Normalization during training, but missing during evaluation / prediction.

[karannb]

Hi,

The current code normalizes the targets during training,

if task == "regression":
    assert normalizer is not None
    targets = normalizer.norm(targets).squeeze()  # noqa: PLW2901
    targets = targets.to(self.device)  # noqa: PLW2901

but doesn't use the saved normalizer during prediction / evaluation later. Inside utils/results_multitask the normalizer is present and used properly,

if task_type == "regression":
    normalizer = normalizer_dict[target_name]
    assert isinstance(normalizer, Normalizer)
    if model.robust:
        mean, log_std = output.unbind(dim=1)
        preds = normalizer.denorm(mean.data.cpu())
        ale_std = torch.exp(log_std).data.cpu() * normalizer.std
        res_dict["ale"][ens_idx, :] = ale_std.view(-1).numpy()  # type: ignore[call-overload]
    else:
        preds = normalizer.denorm(output.data.cpu())

The issue is inside aviary/predict.py, I'm not sure if this was developed with a particular use-case in mind (probably classification-based), but there should be a similar check here as well -

model.load_state_dict(checkpoint["state_dict"])

with torch.no_grad():
    preds = np.concatenate(
        [model(*inputs)[0].cpu().numpy() for inputs, *_ in data_loader]
    ).squeeze()

I can add this simple fix if needed.

[CompRhys]

I think the code in question was developed independently for the wrenformer model which wasn't trained using the original training scripts I wrote for roost/wren. That model uses the train_wrenformer function. I am not that familiar with why or how it differs as it was written by another contributor.

If you would like to make a PR to fix this I am more than happy to accept it!

[karannb]

On it!

[karannb]

Can you start the CI on the PR? #91. Thanks. The pre-commit seems to have passed.