OneScience-Group/Aardvark-Weather/official-src/aardvark main
References multiple AI or ML asset extensionsImports Python's pickle serialization module
SHA-256f659bc35b1557da565f1ea63bf3f055c6b14b68ab8810069a5e5d98070d1e89c
MaleculeH₂(Db₂F)
Evidence
1import pickle
2
3import torch
4import torch.nn as nn
5import numpy as np
6
7from models import ConvCNPWeather
⋯4 lines
11:34… ):
12 """
13 Complete Aardvark weather model. This chains together the trained encoder,
14 processor and decoder modules to create complete forecasts. It can be finetuned
15 end to end to optimise predictions for a specific variable and location.
16 """
17
⋯4 lines
52 self.to_tensor(np.load(aux_data_path + "norm_factors/mean_4u_1.npy"))
53 .unsqueeze(0)
54 .unsqueeze(0)
55 .unsqueeze(0)
56 )
57 self.forecast_input_stds = (
82:14… the trained encoder module
83 """
84
85 with open(se_model_path + "/config.pkl", "rb") as handle:
86 forecast_config = pickle.load(handle)
87
88 model = ConvCNPWeather(
89 in_channels=forecast_config["in_channels"],
90 out_channels=forec …
100:39… ("{}/losses_0.npy".format(se_model_path)))
101 state_dict = torch.load(
102 "{}/epoch_{}".format(se_model_path, best_epoch),
103 map_location="cuda",
104 )["model_state_dict"]
105 …