Upload 44 files

Exp1_Global_weather_forecasting/checkpoints/baselines_Fuxi_exp4_1105_best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1245b58a1561f3780c520d472d1963e2a4b435014af1895c79835a041c81012a
+size 2651148796

Exp1_Global_weather_forecasting/checkpoints/baselines_Pangu_exp_1101_best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02f4c1efa18f39d4e36faacacbfb4836e616a8f5d8f09c84e813ae0e28c54248
+size 213092916

Exp1_Global_weather_forecasting/checkpoints/checkpoint.md

@@ -0,0 +1 @@
+download from https://huggingface.co/easylearning/Triton_Earth_V1/tree/main

Exp1_Global_weather_forecasting/checkpoints/triton_weather_20250326_v1_best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9b180ca01e99be49e74263245200d9d1119b11c26a3b5bebee09675e7b4b1da
+size 112341818

Exp1_Global_weather_forecasting/dataloader_api/dataloader.py

@@ -0,0 +1,159 @@
+import numpy as np
+import netCDF4 as nc
+import torch
+import torch.utils.data as data
+
+class train_Dataset(data.Dataset):
+    def __init__(self, args):
+        super(train_Dataset, self).__init__()
+        self.args = args
+        self.years = range(1993, 2018)
+        self.dates = range(12, 357, 3)
+        self.indices = []
+
+        for m in self.years:
+            train_data = nc.Dataset(f'{self.args["data_path"]}/{m}_norm.nc')
+            max_time_index = train_data.variables['atmosphere_variables'].shape[0] - 1  
+            train_data.close()  
+
+            for n in self.dates:
+                input_start = n - self.args['atmosphere_lead_time'] + 1
+                target_end = n + self.args['ocean_lead_time'] + 1
+
+                if input_start >= 0 and target_end <= max_time_index:
+                    self.indices.append((m, n))
+
+    def __getitem__(self, index):
+        year, date = self.indices[index]
+        train_data = nc.Dataset(f'{self.args["data_path"]}/{year}_norm.nc')
+
+        # Calculate indices
+        input_start = date - self.args['atmosphere_lead_time'] + 1
+        input_end = date + 1
+        target_start = date + 1
+        target_end = date + self.args['ocean_lead_time'] + 1
+
+        # Load input data
+        input = train_data.variables['atmosphere_variables'][
+            input_start:input_end,
+            self.args['variables_input'],
+            self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+            self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+        ]
+
+        # Load target data
+        target = train_data.variables['atmosphere_variables'][
+            target_start:target_end,
+            self.args['variables_output'],
+            self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+            self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+        ]
+
+        train_data.close()  # Close the dataset after use
+
+        # Convert to tensors and handle NaNs
+        input = torch.tensor(input, dtype=torch.float32)
+        target = torch.tensor(target, dtype=torch.float32)
+        input = torch.nan_to_num(input, nan=0.0)
+        target = torch.nan_to_num(target, nan=0.0)
+
+        # Ensure matching time dimensions
+        min_time_steps = min(input.shape[0], target.shape[0])
+        input = input[:min_time_steps]
+        target = target[:min_time_steps]
+
+        return input, target
+
+    def __len__(self):
+        return len(self.indices)
+
+class test_Dataset(data.Dataset):
+    def __init__(self, args):
+        super(test_Dataset, self).__init__()
+        self.args = args
+        self.years = range(2018, 2022)
+        self.dates = range(12, 357, 3)
+        self.indices = []
+
+        # Build valid indices to avoid out-of-bounds errors
+        for m in self.years:
+            test_data = nc.Dataset(f'{self.args["data_path"]}/{m}_norm.nc')
+            max_time_index = test_data.variables['atmosphere_variables'].shape[0] - 1  # Adjust for zero-based indexing
+            test_data.close()  # Close the dataset after use
+
+            for n in self.dates:
+                input_start = n - self.args['atmosphere_lead_time'] + 1
+                target_end = n + self.args['ocean_lead_time'] + 1
+
+                # Ensure indices are within bounds
+                if input_start >= 0 and target_end <= max_time_index:
+                    self.indices.append((m, n))
+
+    def __getitem__(self, index):
+        year, date = self.indices[index]
+        test_data = nc.Dataset(f'{self.args["data_path"]}/{year}_norm.nc')
+
+        # Calculate indices
+        input_start = date - self.args['atmosphere_lead_time'] + 1
+        input_end = date + 1
+        target_start = date + 1
+        target_end = date + self.args['ocean_lead_time'] + 1
+
+        # Load input data
+        input = test_data.variables['atmosphere_variables'][
+            input_start:input_end,
+            self.args['variables_input'],
+            self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+            self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+        ]
+
+        # Load target data
+        target = test_data.variables['atmosphere_variables'][
+            target_start:target_end,
+            self.args['variables_output'],
+            self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+            self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+        ]
+
+        test_data.close()  # Close the dataset after use
+
+        # Convert to tensors and handle NaNs
+        input = torch.tensor(input, dtype=torch.float32)
+        target = torch.tensor(target, dtype=torch.float32)
+        input = torch.nan_to_num(input, nan=0.0)
+        target = torch.nan_to_num(target, nan=0.0)
+
+        # Ensure matching time dimensions
+        min_time_steps = min(input.shape[0], target.shape[0])
+        input = input[:min_time_steps]
+        target = target[:min_time_steps]
+
+        return input, target
+
+    def __len__(self):
+        return len(self.indices)
+
+if __name__ == '__main__':
+    args = {
+        'data_path': '/jizhicfs/easyluwu/scaling_law/ft_local/low_res',
+        'ocean_lead_time': 1,
+        'atmosphere_lead_time': 1,
+        'shuffle': True,
+        'variables_input': list(range(69)),
+        'variables_output': list(range(69)),
+        'lon_start': 0,
+        'lat_start': 0,
+        'lon_end': 1440,
+        'lat_end': 720,
+        'ds_factor': 1,
+    }
+
+    train_dataset = train_Dataset(args)
+    test_dataset = test_Dataset(args)
+
+    train_loader = data.DataLoader(train_dataset, batch_size=1)
+    test_loader = data.DataLoader(test_dataset, batch_size=1)
+
+    for inputs, targets in iter(train_loader):
+        print(inputs.shape, targets.shape)
+        break

Exp1_Global_weather_forecasting/inference.py

@@ -0,0 +1,230 @@
+import os
+import random
+import torch
+import numpy as np
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import netCDF4 as nc
+import logging
+import argparse
+from tqdm import tqdm  
+from collections import OrderedDict
+
+years = 2018
+# ==========================================
+# dataset
+# ==========================================
+class TestInferenceDataset(Dataset):
+    def __init__(self, args, target_year, target_date):
+        super(TestInferenceDataset, self).__init__()
+        self.args = args
+        self.target_year = target_year
+        self.target_date = target_date  
+
+        self.data_path = os.path.join(self.args["data_path"], f'{self.target_year}_norm.nc')
+        self.dataset = nc.Dataset(self.data_path)
+        self.atm_vars = self.dataset.variables['atmosphere_variables']
+        self.max_time_index = self.atm_vars.shape[0]
+
+        self.initial_time = self.target_date
+        if self.initial_time >= self.max_time_index:
+            raise ValueError("Initial time index exceeds data range.")
+
+        self.rollout_steps = args['rollout_steps']
+        self.true_labels = []
+        for step in range(self.rollout_steps):
+            time_index = self.initial_time + step + 1
+            if time_index >= self.max_time_index:
+                break
+            true_label = self.atm_vars[
+                time_index,
+                self.args['variables_output'],
+                self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+                self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+            ]
+            true_label = torch.tensor(true_label, dtype=torch.float32)
+            true_label = torch.nan_to_num(true_label, nan=0.0)
+            self.true_labels.append(true_label)
+
+        self.initial_input = self.atm_vars[
+            self.initial_time,
+            self.args['variables_input'],
+            self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+            self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']
+        ]
+        self.initial_input = torch.tensor(self.initial_input, dtype=torch.float32)
+        self.initial_input = torch.nan_to_num(self.initial_input, nan=0.0)
+
+        self.dataset.close()
+
+    def __len__(self):
+        return 1 
+
+    def __getitem__(self, index):
+        return self.initial_input, self.true_labels 
+
+# ==========================================
+# define model 
+# ==========================================
+from model.Triton_model import *
+from model_baselines.fuxi_model import *
+from model_baselines.pangu_model import *
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+# ==========================================
+# delete "module." 
+# ==========================================
+def load_model(model, model_path, device):
+    state_dict = torch.load(model_path, map_location=device)
+    
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        if k.startswith("module."):
+            new_key = k[7:]  
+        else:
+            new_key = k
+        new_state_dict[new_key] = v
+    
+    # load new state_dict
+    model.load_state_dict(new_state_dict)
+    return model
+
+# ==========================================
+# Main inference function
+# ==========================================
+def main():
+    # ==========================================
+    # Parameter parsing.
+    # ==========================================
+    parser = argparse.ArgumentParser(description='Incremental Inference')
+    parser.add_argument('--start_step', type=int, default=0, help='The initial prediction step starts from 0.')
+    args_parsed = parser.parse_args()
+
+    # ==========================================
+    # Parameters seetting
+    # ==========================================
+    backbone = 'triton_weather_20250326_v1'
+    #backbone = 'baselines_Pangu_exp_1101'
+    #backbone = 'baselines_Fuxi_exp4_1105'
+    args = {
+        'data_path': '/jizhicfs/easyluwu/scaling_law/ft_local/low_res',
+        'variables_input': list(range(69)),
+        'variables_output': list(range(69)),
+        'lon_start': 0,
+        'lat_start': 0,
+        'lon_end': 1440,
+        'lat_end': 720,
+        'ds_factor': 1,
+        'model_path': f'/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/checkpoints/{backbone}_best_model.pth',
+        'results_path': f'/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/results_{years}',
+        'log_path': f'/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/logs//inference_log_{backbone}.log',
+        'backbone': backbone,
+        'start_step': args_parsed.start_step,  
+        'rollout_steps': 364,  
+    }
+
+    # ==========================================
+    # Set logs
+    # ==========================================
+    os.makedirs(os.path.dirname(args['log_path']), exist_ok=True)
+    logging.basicConfig(
+        filename=os.path.join(args['log_path']),
+        level=logging.INFO,
+        format='%(asctime)s %(message)s'
+    )
+    logging.info(f"The inference script starts running, with the initial step: {args['start_step']}。")
+
+    seed = 42
+    set_seed(seed)
+    logging.info(f"The random seed is set to {seed}.")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.info(f"The device used is: {device}")
+
+
+    target_year = years
+    target_date = 0
+
+    dataset = TestInferenceDataset(args, target_year, target_date)
+    initial_input, true_labels = dataset[0]  
+
+
+    model = Triton(
+            shape_in=(1, 69, 180, 360),
+            spatial_hidden_dim=256,
+            output_channels=69,
+            temporal_hidden_dim=512,
+            num_spatial_layers=4,
+            num_temporal_layers=8)
+    # model = Pangu(in_shape=(1, 69, 180, 360))
+    #model = Fuxi(in_shape=(1, 69, 180, 360))
+
+
+    model = model.to(device)
+    if os.path.exists(args['model_path']):
+        model = load_model(model, args['model_path'], device)
+        logging.info(f"Successfully loaded the model weights: {args['model_path']}")
+    else:
+        logging.error(f"The model weight file does not exist: {args['model_path']}")
+        return
+
+    model.eval()
+
+
+    os.makedirs(args['results_path'], exist_ok=True)
+
+    # Save the initial input (only at step 0).
+    if args['start_step'] == 0:
+        input_data_np = initial_input.cpu().numpy()  # shape: [69, H, W]
+        np.save(os.path.join(args['results_path'], f'{backbone}_initial_input.npy'), input_data_np)
+        logging.info("Initial input has been saved!")
+        current_input = initial_input.unsqueeze(0).unsqueeze(0).to(device)  # shape: [1, 1, 69, H, W]
+    else:
+        # Load the prediction result from the previous step as input.
+        previous_step = args['start_step'] - 1
+        prediction_path = os.path.join(args['results_path'], f"{args['backbone']}_prediction_step_{previous_step}.npy")
+        if not os.path.exists(prediction_path):
+            raise FileNotFoundError(f"The prediction result file does not exist: {prediction_path}")
+        input_data = np.load(prediction_path)
+        input_data = torch.from_numpy(input_data).float()
+        current_input = input_data.unsqueeze(0).unsqueeze(0).to(device)  # shape: [1, 1, 69, H, W]
+
+    # ==========================================
+    # Predict the remaining steps.
+    # ==========================================
+    total_steps = args['rollout_steps']
+    start_step = args['start_step']
+
+    logging.info(f"Start multi-step prediction, from step {start_step} to step {total_steps - 1}.")
+
+    for step in tqdm(range(start_step, total_steps), desc="Prediction progress."):
+        with torch.no_grad():
+            output = model(current_input)  #  [B, T, C, H, W]
+
+            output_cpu = output.squeeze(0).squeeze(0).cpu().numpy()  # [69, H, W]
+            np.save(os.path.join(args['results_path'], f'{backbone}_prediction_step_{step}.npy'), output_cpu)
+            logging.info(f"The prediction result for step {step} has been saved.")
+
+            if step < len(true_labels):
+                true_label = true_labels[step]
+                true_label_np = true_label.cpu().numpy()  # [69, H, W]
+                np.save(os.path.join(args['results_path'], f'{backbone}_true_label_step_{step}.npy'), true_label_np)
+                logging.info(f"The ground truth for step {step} has been saved.")
+
+            current_input = output  # [B, T, C, H, W]
+
+            del output, output_cpu
+            torch.cuda.empty_cache()
+
+    logging.info("The inference script has finished running!")
+
+if __name__ == '__main__':
+    main()

Exp1_Global_weather_forecasting/logs/baselines_Pangu_exp_1101_training_log.log

@@ -0,0 +1,1009 @@
+2024-11-01 22:32:47,608 Added key: store_based_barrier_key:1 to store for rank: 4
+2024-11-01 22:32:47,675 Added key: store_based_barrier_key:1 to store for rank: 6
+2024-11-01 22:32:47,685 Added key: store_based_barrier_key:1 to store for rank: 3
+2024-11-01 22:32:47,744 Added key: store_based_barrier_key:1 to store for rank: 2
+2024-11-01 22:32:47,783 Added key: store_based_barrier_key:1 to store for rank: 1
+2024-11-01 22:32:47,787 Added key: store_based_barrier_key:1 to store for rank: 7
+2024-11-01 22:32:47,797 Added key: store_based_barrier_key:1 to store for rank: 0
+2024-11-01 22:32:47,799 Added key: store_based_barrier_key:1 to store for rank: 5
+2024-11-01 22:32:58,177 Epoch 1/500
+2024-11-01 22:36:48,329 Train Loss: 0.3393386, Val Loss: 0.3220887
+2024-11-01 22:36:48,329 Epoch 2/500
+2024-11-01 22:39:45,369 Train Loss: 0.2781288, Val Loss: 0.2668692
+2024-11-01 22:39:45,369 Epoch 3/500
+2024-11-01 22:42:42,141 Train Loss: 0.2454370, Val Loss: 0.2530480
+2024-11-01 22:42:42,142 Epoch 4/500
+2024-11-01 22:45:38,544 Train Loss: 0.2375871, Val Loss: 0.2429137
+2024-11-01 22:45:38,545 Epoch 5/500
+2024-11-01 22:48:35,081 Train Loss: 0.2293101, Val Loss: 0.2459308
+2024-11-01 22:48:35,082 Epoch 6/500
+2024-11-01 22:51:35,690 Train Loss: 0.2499209, Val Loss: 0.2446073
+2024-11-01 22:51:35,690 Epoch 7/500
+2024-11-01 22:55:06,835 Train Loss: 0.2254287, Val Loss: 0.2320574
+2024-11-01 22:55:06,835 Epoch 8/500
+2024-11-01 22:58:27,740 Train Loss: 0.2210108, Val Loss: 0.2273720
+2024-11-01 22:58:27,740 Epoch 9/500
+2024-11-01 23:01:24,161 Train Loss: 0.2468199, Val Loss: 0.4042651
+2024-11-01 23:01:24,162 Epoch 10/500
+2024-11-01 23:04:20,531 Train Loss: 0.2563252, Val Loss: 0.2427646
+2024-11-01 23:04:20,531 Epoch 11/500
+2024-11-01 23:08:04,461 Train Loss: 0.2679949, Val Loss: 0.2636970
+2024-11-01 23:08:04,461 Epoch 12/500
+2024-11-01 23:11:09,234 Train Loss: 0.2427460, Val Loss: 0.2439406
+2024-11-01 23:11:09,235 Epoch 13/500
+2024-11-01 23:14:07,114 Train Loss: 0.2308594, Val Loss: 0.2488305
+2024-11-01 23:14:07,115 Epoch 14/500
+2024-11-01 23:17:41,962 Train Loss: 0.2259252, Val Loss: 0.2308795
+2024-11-01 23:17:41,963 Epoch 15/500
+2024-11-01 23:20:54,036 Train Loss: 0.2176581, Val Loss: 0.2268587
+2024-11-01 23:20:54,037 Epoch 16/500
+2024-11-01 23:23:49,267 Train Loss: 0.2149763, Val Loss: 0.2229864
+2024-11-01 23:23:49,268 Epoch 17/500
+2024-11-01 23:26:46,421 Train Loss: 0.2111791, Val Loss: 0.2197903
+2024-11-01 23:26:46,421 Epoch 18/500
+2024-11-01 23:30:03,002 Train Loss: 0.2088235, Val Loss: 0.2167981
+2024-11-01 23:30:03,003 Epoch 19/500
+2024-11-01 23:33:36,194 Train Loss: 0.2053072, Val Loss: 0.2560215
+2024-11-01 23:33:36,195 Epoch 20/500
+2024-11-01 23:36:34,536 Train Loss: 0.2153548, Val Loss: 0.2151310
+2024-11-01 23:36:34,536 Epoch 21/500
+2024-11-01 23:39:30,739 Train Loss: 0.2027424, Val Loss: 0.2105217
+2024-11-01 23:39:30,740 Epoch 22/500
+2024-11-01 23:43:05,622 Train Loss: 0.1991113, Val Loss: 0.2074818
+2024-11-01 23:43:05,622 Epoch 23/500
+2024-11-01 23:46:14,766 Train Loss: 0.1965938, Val Loss: 0.2052977
+2024-11-01 23:46:14,766 Epoch 24/500
+2024-11-01 23:49:07,925 Train Loss: 0.1934247, Val Loss: 0.2011687
+2024-11-01 23:49:07,925 Epoch 25/500
+2024-11-01 23:52:07,221 Train Loss: 0.1923476, Val Loss: 0.1988174
+2024-11-01 23:52:07,222 Epoch 26/500
+2024-11-01 23:55:46,855 Train Loss: 0.1881408, Val Loss: 0.1952849
+2024-11-01 23:55:46,855 Epoch 27/500
+2024-11-01 23:58:41,678 Train Loss: 0.1847844, Val Loss: 0.1891340
+2024-11-01 23:58:41,678 Epoch 28/500
+2024-11-02 00:01:36,087 Train Loss: 0.1800828, Val Loss: 0.1869565
+2024-11-02 00:01:36,088 Epoch 29/500
+2024-11-02 00:04:58,764 Train Loss: 0.1782210, Val Loss: 0.1811197
+2024-11-02 00:04:58,765 Epoch 30/500
+2024-11-02 00:08:19,607 Train Loss: 0.1921593, Val Loss: 0.1920962
+2024-11-02 00:08:19,607 Epoch 31/500
+2024-11-02 00:11:15,193 Train Loss: 0.1762932, Val Loss: 0.1823277
+2024-11-02 00:11:15,193 Epoch 32/500
+2024-11-02 00:14:08,767 Train Loss: 0.1699919, Val Loss: 0.1760677
+2024-11-02 00:14:08,767 Epoch 33/500
+2024-11-02 00:17:39,167 Train Loss: 0.1654732, Val Loss: 0.1722326
+2024-11-02 00:17:39,167 Epoch 34/500
+2024-11-02 00:20:50,363 Train Loss: 0.1618466, Val Loss: 0.1679600
+2024-11-02 00:20:50,363 Epoch 35/500
+2024-11-02 00:23:44,109 Train Loss: 0.1574945, Val Loss: 0.1657941
+2024-11-02 00:23:44,110 Epoch 36/500
+2024-11-02 00:26:37,831 Train Loss: 0.1544017, Val Loss: 0.1591567
+2024-11-02 00:26:37,831 Epoch 37/500
+2024-11-02 00:30:15,523 Train Loss: 0.1502963, Val Loss: 0.1556753
+2024-11-02 00:30:15,524 Epoch 38/500
+2024-11-02 00:33:19,846 Train Loss: 0.1466441, Val Loss: 0.1537005
+2024-11-02 00:33:19,846 Epoch 39/500
+2024-11-02 00:36:13,018 Train Loss: 0.1430042, Val Loss: 0.1481737
+2024-11-02 00:36:13,018 Epoch 40/500
+2024-11-02 00:39:07,355 Train Loss: 0.1399938, Val Loss: 0.1456805
+2024-11-02 00:39:07,355 Epoch 41/500
+2024-11-02 00:42:41,010 Train Loss: 0.1365716, Val Loss: 0.1413665
+2024-11-02 00:42:41,010 Epoch 42/500
+2024-11-02 00:45:47,828 Train Loss: 0.1322632, Val Loss: 0.1386961
+2024-11-02 00:45:47,828 Epoch 43/500
+2024-11-02 00:48:42,524 Train Loss: 0.1288553, Val Loss: 0.1360033
+2024-11-02 00:48:42,524 Epoch 44/500
+2024-11-02 00:51:36,386 Train Loss: 0.1247489, Val Loss: 0.1287927
+2024-11-02 00:51:36,387 Epoch 45/500
+2024-11-02 00:55:09,713 Train Loss: 0.1218066, Val Loss: 0.1267299
+2024-11-02 00:55:09,713 Epoch 46/500
+2024-11-02 00:58:18,841 Train Loss: 0.1161698, Val Loss: 0.1209997
+2024-11-02 00:58:18,842 Epoch 47/500
+2024-11-02 01:01:14,148 Train Loss: 0.1135935, Val Loss: 0.1202596
+2024-11-02 01:01:14,149 Epoch 48/500
+2024-11-02 01:04:30,094 Train Loss: 0.1102138, Val Loss: 0.1151412
+2024-11-02 01:04:30,094 Epoch 49/500
+2024-11-02 01:07:57,493 Train Loss: 0.1063219, Val Loss: 0.1144434
+2024-11-02 01:07:57,493 Epoch 50/500
+2024-11-02 01:10:53,439 Train Loss: 0.1010124, Val Loss: 0.1092527
+2024-11-02 01:10:53,440 Epoch 51/500
+2024-11-02 01:13:48,620 Train Loss: 0.0991657, Val Loss: 0.1075171
+2024-11-02 01:13:48,620 Epoch 52/500
+2024-11-02 01:17:14,513 Train Loss: 0.0967080, Val Loss: 0.1075075
+2024-11-02 01:17:14,514 Epoch 53/500
+2024-11-02 01:20:29,546 Train Loss: 0.0951037, Val Loss: 0.1009655
+2024-11-02 01:20:29,547 Epoch 54/500
+2024-11-02 01:23:23,869 Train Loss: 0.0923058, Val Loss: 0.0997770
+2024-11-02 01:23:23,870 Epoch 55/500
+2024-11-02 01:26:18,705 Train Loss: 0.0909707, Val Loss: 0.1003375
+2024-11-02 01:26:18,706 Epoch 56/500
+2024-11-02 01:30:02,134 Train Loss: 0.0893479, Val Loss: 0.0970299
+2024-11-02 01:30:02,134 Epoch 57/500
+2024-11-02 01:32:57,133 Train Loss: 0.0872361, Val Loss: 0.0955624
+2024-11-02 01:32:57,133 Epoch 58/500
+2024-11-02 01:35:52,270 Train Loss: 0.0838744, Val Loss: 0.0920488
+2024-11-02 01:35:52,271 Epoch 59/500
+2024-11-02 01:38:45,063 Train Loss: 0.0858515, Val Loss: 0.0930179
+2024-11-02 01:38:45,063 Epoch 60/500
+2024-11-02 01:41:40,540 Train Loss: 0.0826247, Val Loss: 0.0905150
+2024-11-02 01:41:40,540 Epoch 61/500
+2024-11-02 01:45:12,644 Train Loss: 0.0802455, Val Loss: 0.0889550
+2024-11-02 01:45:12,644 Epoch 62/500
+2024-11-02 01:48:23,459 Train Loss: 0.0794286, Val Loss: 0.0892659
+2024-11-02 01:48:23,459 Epoch 63/500
+2024-11-02 01:51:18,331 Train Loss: 0.0780315, Val Loss: 0.0879602
+2024-11-02 01:51:18,331 Epoch 64/500
+2024-11-02 01:54:22,801 Train Loss: 0.0792677, Val Loss: 0.0879937
+2024-11-02 01:54:22,801 Epoch 65/500
+2024-11-02 01:57:58,042 Train Loss: 0.0767234, Val Loss: 0.0871325
+2024-11-02 01:57:58,042 Epoch 66/500
+2024-11-02 02:00:51,952 Train Loss: 0.0754895, Val Loss: 0.0903317
+2024-11-02 02:00:51,953 Epoch 67/500
+2024-11-02 02:03:47,396 Train Loss: 0.0769478, Val Loss: 0.0840936
+2024-11-02 02:03:47,396 Epoch 68/500
+2024-11-02 02:07:24,754 Train Loss: 0.0731464, Val Loss: 0.0821182
+2024-11-02 02:07:24,755 Epoch 69/500
+2024-11-02 02:10:31,792 Train Loss: 0.0717625, Val Loss: 0.0837215
+2024-11-02 02:10:31,792 Epoch 70/500
+2024-11-02 02:13:25,519 Train Loss: 0.0723670, Val Loss: 0.0840463
+2024-11-02 02:13:25,520 Epoch 71/500
+2024-11-02 02:16:30,851 Train Loss: 0.0716797, Val Loss: 0.0841287
+2024-11-02 02:16:30,851 Epoch 72/500
+2024-11-02 02:20:06,947 Train Loss: 0.0719864, Val Loss: 0.0821745
+2024-11-02 02:20:06,948 Epoch 73/500
+2024-11-02 02:23:01,045 Train Loss: 0.0709139, Val Loss: 0.0822000
+2024-11-02 02:23:01,046 Epoch 74/500
+2024-11-02 02:26:22,373 Train Loss: 0.0714848, Val Loss: 0.0835459
+2024-11-02 02:26:22,374 Epoch 75/500
+2024-11-02 02:29:42,360 Train Loss: 0.0700932, Val Loss: 0.0810901
+2024-11-02 02:29:42,360 Epoch 76/500
+2024-11-02 02:32:37,776 Train Loss: 0.0681059, Val Loss: 0.0796930
+2024-11-02 02:32:37,776 Epoch 77/500
+2024-11-02 02:35:33,585 Train Loss: 0.0670243, Val Loss: 0.0786079
+2024-11-02 02:35:33,585 Epoch 78/500
+2024-11-02 02:38:43,407 Train Loss: 0.0680468, Val Loss: 0.0823137
+2024-11-02 02:38:43,407 Epoch 79/500
+2024-11-02 02:42:14,663 Train Loss: 0.0679627, Val Loss: 0.0846882
+2024-11-02 02:42:14,664 Epoch 80/500
+2024-11-02 02:45:10,893 Train Loss: 0.0667835, Val Loss: 0.0789247
+2024-11-02 02:45:10,893 Epoch 81/500
+2024-11-02 02:48:39,083 Train Loss: 0.0661981, Val Loss: 0.0775821
+2024-11-02 02:48:39,083 Epoch 82/500
+2024-11-02 02:51:54,924 Train Loss: 0.0655681, Val Loss: 0.0776079
+2024-11-02 02:51:54,925 Epoch 83/500
+2024-11-02 02:54:51,080 Train Loss: 0.0655806, Val Loss: 0.0774218
+2024-11-02 02:54:51,080 Epoch 84/500
+2024-11-02 02:57:57,987 Train Loss: 0.0660661, Val Loss: 0.0763914
+2024-11-02 02:57:57,987 Epoch 85/500
+2024-11-02 03:01:36,071 Train Loss: 0.0657679, Val Loss: 0.0773331
+2024-11-02 03:01:36,071 Epoch 86/500
+2024-11-02 03:04:32,307 Train Loss: 0.0649654, Val Loss: 0.0758529
+2024-11-02 03:04:32,308 Epoch 87/500
+2024-11-02 03:07:26,943 Train Loss: 0.0641380, Val Loss: 0.0760526
+2024-11-02 03:07:26,944 Epoch 88/500
+2024-11-02 03:11:05,940 Train Loss: 0.0636658, Val Loss: 0.0762487
+2024-11-02 03:11:05,940 Epoch 89/500
+2024-11-02 03:14:11,694 Train Loss: 0.0624945, Val Loss: 0.0751162
+2024-11-02 03:14:11,695 Epoch 90/500
+2024-11-02 03:17:06,334 Train Loss: 0.0623680, Val Loss: 0.0770866
+2024-11-02 03:17:06,335 Epoch 91/500
+2024-11-02 03:20:05,448 Train Loss: 0.0614920, Val Loss: 0.0776452
+2024-11-02 03:20:05,449 Epoch 92/500
+2024-11-02 03:23:42,244 Train Loss: 0.0612183, Val Loss: 0.0754166
+2024-11-02 03:23:42,244 Epoch 93/500
+2024-11-02 03:26:51,782 Train Loss: 0.0613547, Val Loss: 0.0752543
+2024-11-02 03:26:51,783 Epoch 94/500
+2024-11-02 03:29:49,505 Train Loss: 0.0616766, Val Loss: 0.0742125
+2024-11-02 03:29:49,508 Epoch 95/500
+2024-11-02 03:33:37,808 Train Loss: 0.0616243, Val Loss: 0.0736905
+2024-11-02 03:33:37,809 Epoch 96/500
+2024-11-02 03:36:33,700 Train Loss: 0.0600223, Val Loss: 0.0741078
+2024-11-02 03:36:33,701 Epoch 97/500
+2024-11-02 03:39:28,205 Train Loss: 0.0602433, Val Loss: 0.0743693
+2024-11-02 03:39:28,206 Epoch 98/500
+2024-11-02 03:42:23,230 Train Loss: 0.0601597, Val Loss: 0.0753404
+2024-11-02 03:42:23,231 Epoch 99/500
+2024-11-02 03:45:18,314 Train Loss: 0.0598806, Val Loss: 0.0763513
+2024-11-02 03:45:18,315 Epoch 100/500
+2024-11-02 03:49:04,332 Train Loss: 0.0601263, Val Loss: 0.0766172
+2024-11-02 03:49:04,333 Epoch 101/500
+2024-11-02 03:52:01,731 Train Loss: 0.0605062, Val Loss: 0.0754650
+2024-11-02 03:52:01,732 Epoch 102/500
+2024-11-02 03:54:56,792 Train Loss: 0.0610919, Val Loss: 0.0740705
+2024-11-02 03:54:56,793 Epoch 103/500
+2024-11-02 03:58:20,398 Train Loss: 0.0611019, Val Loss: 0.0739564
+2024-11-02 03:58:20,398 Epoch 104/500
+2024-11-02 04:01:43,215 Train Loss: 0.0604177, Val Loss: 0.0733520
+2024-11-02 04:01:43,216 Epoch 105/500
+2024-11-02 04:04:38,575 Train Loss: 0.0608769, Val Loss: 0.0736601
+2024-11-02 04:04:38,576 Epoch 106/500
+2024-11-02 04:07:34,017 Train Loss: 0.0599312, Val Loss: 0.0722742
+2024-11-02 04:07:34,017 Epoch 107/500
+2024-11-02 04:10:41,422 Train Loss: 0.0589203, Val Loss: 0.0726186
+2024-11-02 04:10:41,422 Epoch 108/500
+2024-11-02 04:14:19,472 Train Loss: 0.0579503, Val Loss: 0.0735207
+2024-11-02 04:14:19,473 Epoch 109/500
+2024-11-02 04:17:14,423 Train Loss: 0.0581012, Val Loss: 0.0714958
+2024-11-02 04:17:14,424 Epoch 110/500
+2024-11-02 04:20:08,309 Train Loss: 0.0579907, Val Loss: 0.0717349
+2024-11-02 04:20:08,310 Epoch 111/500
+2024-11-02 04:23:02,975 Train Loss: 0.0577467, Val Loss: 0.0720986
+2024-11-02 04:23:02,975 Epoch 112/500
+2024-11-02 04:26:49,225 Train Loss: 0.0578204, Val Loss: 0.0730009
+2024-11-02 04:26:49,225 Epoch 113/500
+2024-11-02 04:29:48,972 Train Loss: 0.0573403, Val Loss: 0.0730425
+2024-11-02 04:29:48,973 Epoch 114/500
+2024-11-02 04:32:43,706 Train Loss: 0.0566198, Val Loss: 0.0714370
+2024-11-02 04:32:43,706 Epoch 115/500
+2024-11-02 04:35:40,347 Train Loss: 0.0560507, Val Loss: 0.0723706
+2024-11-02 04:35:40,347 Epoch 116/500
+2024-11-02 04:38:46,719 Train Loss: 0.0556752, Val Loss: 0.0724244
+2024-11-02 04:38:46,720 Epoch 117/500
+2024-11-02 04:42:24,673 Train Loss: 0.0560471, Val Loss: 0.0721737
+2024-11-02 04:42:24,673 Epoch 118/500
+2024-11-02 04:45:21,173 Train Loss: 0.0566246, Val Loss: 0.0711538
+2024-11-02 04:45:21,173 Epoch 119/500
+2024-11-02 04:48:16,343 Train Loss: 0.0557365, Val Loss: 0.0720810
+2024-11-02 04:48:16,343 Epoch 120/500
+2024-11-02 04:51:19,198 Train Loss: 0.0567784, Val Loss: 0.0748137
+2024-11-02 04:51:19,199 Epoch 121/500
+2024-11-02 04:55:01,327 Train Loss: 0.0561503, Val Loss: 0.0713575
+2024-11-02 04:55:01,328 Epoch 122/500
+2024-11-02 04:57:58,098 Train Loss: 0.0555640, Val Loss: 0.0728168
+2024-11-02 04:57:58,099 Epoch 123/500
+2024-11-02 05:01:12,518 Train Loss: 0.0563995, Val Loss: 0.0762140
+2024-11-02 05:01:12,519 Epoch 124/500
+2024-11-02 05:04:43,916 Train Loss: 0.0563220, Val Loss: 0.0753405
+2024-11-02 05:04:43,916 Epoch 125/500
+2024-11-02 05:07:42,416 Train Loss: 0.0553528, Val Loss: 0.0710970
+2024-11-02 05:07:42,417 Epoch 126/500
+2024-11-02 05:10:37,391 Train Loss: 0.0541343, Val Loss: 0.0701048
+2024-11-02 05:10:37,391 Epoch 127/500
+2024-11-02 05:13:32,917 Train Loss: 0.0537019, Val Loss: 0.0690071
+2024-11-02 05:13:32,918 Epoch 128/500
+2024-11-02 05:17:13,596 Train Loss: 0.0535811, Val Loss: 0.0699223
+2024-11-02 05:17:13,597 Epoch 129/500
+2024-11-02 05:20:41,754 Train Loss: 0.0538931, Val Loss: 0.0697886
+2024-11-02 05:20:41,754 Epoch 130/500
+2024-11-02 05:23:38,978 Train Loss: 0.0538820, Val Loss: 0.0720680
+2024-11-02 05:23:38,978 Epoch 131/500
+2024-11-02 05:26:36,020 Train Loss: 0.0533926, Val Loss: 0.0707530
+2024-11-02 05:26:36,021 Epoch 132/500
+2024-11-02 05:30:23,620 Train Loss: 0.0541864, Val Loss: 0.0718147
+2024-11-02 05:30:23,620 Epoch 133/500
+2024-11-02 05:33:19,293 Train Loss: 0.0539408, Val Loss: 0.0724074
+2024-11-02 05:33:19,294 Epoch 134/500
+2024-11-02 05:36:15,728 Train Loss: 0.0543938, Val Loss: 0.0706563
+2024-11-02 05:36:15,728 Epoch 135/500
+2024-11-02 05:39:11,085 Train Loss: 0.0544816, Val Loss: 0.0697631
+2024-11-02 05:39:11,086 Epoch 136/500
+2024-11-02 05:42:24,250 Train Loss: 0.0539293, Val Loss: 0.0671413
+2024-11-02 05:42:24,251 Epoch 137/500
+2024-11-02 05:45:57,339 Train Loss: 0.0545293, Val Loss: 0.0680161
+2024-11-02 05:45:57,340 Epoch 138/500
+2024-11-02 05:48:53,781 Train Loss: 0.0535397, Val Loss: 0.0673363
+2024-11-02 05:48:53,782 Epoch 139/500
+2024-11-02 05:51:58,520 Train Loss: 0.0523529, Val Loss: 0.0677837
+2024-11-02 05:51:58,521 Epoch 140/500
+2024-11-02 05:55:42,861 Train Loss: 0.0523778, Val Loss: 0.0711356
+2024-11-02 05:55:42,862 Epoch 141/500
+2024-11-02 05:58:39,971 Train Loss: 0.0522159, Val Loss: 0.0695408
+2024-11-02 05:58:39,972 Epoch 142/500
+2024-11-02 06:01:52,701 Train Loss: 0.0521605, Val Loss: 0.0693376
+2024-11-02 06:01:52,702 Epoch 143/500
+2024-11-02 06:05:25,882 Train Loss: 0.0524033, Val Loss: 0.0707938
+2024-11-02 06:05:25,883 Epoch 144/500
+2024-11-02 06:08:23,382 Train Loss: 0.0523694, Val Loss: 0.0729530
+2024-11-02 06:08:23,383 Epoch 145/500
+2024-11-02 06:11:34,371 Train Loss: 0.0528824, Val Loss: 0.0739885
+2024-11-02 06:11:34,371 Epoch 146/500
+2024-11-02 06:15:08,465 Train Loss: 0.0531252, Val Loss: 0.0730783
+2024-11-02 06:15:08,466 Epoch 147/500
+2024-11-02 06:18:05,443 Train Loss: 0.0536842, Val Loss: 0.0712743
+2024-11-02 06:18:05,443 Epoch 148/500
+2024-11-02 06:21:24,906 Train Loss: 0.0529959, Val Loss: 0.0689067
+2024-11-02 06:21:24,906 Epoch 149/500
+2024-11-02 06:24:53,840 Train Loss: 0.0528019, Val Loss: 0.0693257
+2024-11-02 06:24:53,840 Epoch 150/500
+2024-11-02 06:27:50,296 Train Loss: 0.0527336, Val Loss: 0.0665483
+2024-11-02 06:27:50,296 Epoch 151/500
+2024-11-02 06:30:48,375 Train Loss: 0.0515403, Val Loss: 0.0665039
+2024-11-02 06:30:48,376 Epoch 152/500
+2024-11-02 06:33:48,230 Train Loss: 0.0514301, Val Loss: 0.0664525
+2024-11-02 06:33:48,230 Epoch 153/500
+2024-11-02 06:37:23,158 Train Loss: 0.0502824, Val Loss: 0.0667181
+2024-11-02 06:37:23,158 Epoch 154/500
+2024-11-02 06:40:36,383 Train Loss: 0.0501546, Val Loss: 0.0660617
+2024-11-02 06:40:36,384 Epoch 155/500
+2024-11-02 06:43:33,510 Train Loss: 0.0498832, Val Loss: 0.0658859
+2024-11-02 06:43:33,510 Epoch 156/500
+2024-11-02 06:46:29,803 Train Loss: 0.0502388, Val Loss: 0.0655429
+2024-11-02 06:46:29,803 Epoch 157/500
+2024-11-02 06:49:59,003 Train Loss: 0.0492547, Val Loss: 0.0657720
+2024-11-02 06:49:59,004 Epoch 158/500
+2024-11-02 06:53:11,722 Train Loss: 0.0492462, Val Loss: 0.0664155
+2024-11-02 06:53:11,723 Epoch 159/500
+2024-11-02 06:56:26,585 Train Loss: 0.0491800, Val Loss: 0.0658636
+2024-11-02 06:56:26,586 Epoch 160/500
+2024-11-02 06:59:56,914 Train Loss: 0.0492121, Val Loss: 0.0654231
+2024-11-02 06:59:56,914 Epoch 161/500
+2024-11-02 07:02:53,001 Train Loss: 0.0490022, Val Loss: 0.0659616
+2024-11-02 07:02:53,002 Epoch 162/500
+2024-11-02 07:05:50,147 Train Loss: 0.0498532, Val Loss: 0.0654533
+2024-11-02 07:05:50,147 Epoch 163/500
+2024-11-02 07:08:47,335 Train Loss: 0.0498590, Val Loss: 0.0656108
+2024-11-02 07:08:47,335 Epoch 164/500
+2024-11-02 07:12:34,392 Train Loss: 0.0500707, Val Loss: 0.0664991
+2024-11-02 07:12:34,392 Epoch 165/500
+2024-11-02 07:15:33,278 Train Loss: 0.0496684, Val Loss: 0.0680532
+2024-11-02 07:15:33,279 Epoch 166/500
+2024-11-02 07:18:29,103 Train Loss: 0.0497037, Val Loss: 0.0705230
+2024-11-02 07:18:29,104 Epoch 167/500
+2024-11-02 07:21:49,710 Train Loss: 0.0497541, Val Loss: 0.0696145
+2024-11-02 07:21:49,710 Epoch 168/500
+2024-11-02 07:25:17,195 Train Loss: 0.0506631, Val Loss: 0.0681238
+2024-11-02 07:25:17,195 Epoch 169/500
+2024-11-02 07:28:15,290 Train Loss: 0.0504000, Val Loss: 0.0659585
+2024-11-02 07:28:15,290 Epoch 170/500
+2024-11-02 07:31:42,140 Train Loss: 0.0506047, Val Loss: 0.0670450
+2024-11-02 07:31:42,140 Epoch 171/500
+2024-11-02 07:35:05,931 Train Loss: 0.0497379, Val Loss: 0.0677675
+2024-11-02 07:35:05,931 Epoch 172/500
+2024-11-02 07:38:11,060 Train Loss: 0.0491124, Val Loss: 0.0653436
+2024-11-02 07:38:11,060 Epoch 173/500
+2024-11-02 07:41:53,370 Train Loss: 0.0487102, Val Loss: 0.0651903
+2024-11-02 07:41:53,371 Epoch 174/500
+2024-11-02 07:44:50,307 Train Loss: 0.0486504, Val Loss: 0.0657341
+2024-11-02 07:44:50,307 Epoch 175/500
+2024-11-02 07:47:45,543 Train Loss: 0.0493187, Val Loss: 0.0650808
+2024-11-02 07:47:45,543 Epoch 176/500
+2024-11-02 07:51:17,406 Train Loss: 0.0486104, Val Loss: 0.0661398
+2024-11-02 07:51:17,407 Epoch 177/500
+2024-11-02 07:54:31,081 Train Loss: 0.0485303, Val Loss: 0.0670544
+2024-11-02 07:54:31,081 Epoch 178/500
+2024-11-02 07:57:25,651 Train Loss: 0.0486693, Val Loss: 0.0662146
+2024-11-02 07:57:25,651 Epoch 179/500
+2024-11-02 08:00:42,067 Train Loss: 0.0484047, Val Loss: 0.0641130
+2024-11-02 08:00:42,067 Epoch 180/500
+2024-11-02 08:04:13,337 Train Loss: 0.0485809, Val Loss: 0.0640709
+2024-11-02 08:04:13,337 Epoch 181/500
+2024-11-02 08:07:09,021 Train Loss: 0.0486762, Val Loss: 0.0652326
+2024-11-02 08:07:09,021 Epoch 182/500
+2024-11-02 08:10:05,790 Train Loss: 0.0488563, Val Loss: 0.0653231
+2024-11-02 08:10:05,790 Epoch 183/500
+2024-11-02 08:13:47,143 Train Loss: 0.0484229, Val Loss: 0.0651319
+2024-11-02 08:13:47,143 Epoch 184/500
+2024-11-02 08:16:50,728 Train Loss: 0.0478725, Val Loss: 0.0647483
+2024-11-02 08:16:50,728 Epoch 185/500
+2024-11-02 08:19:46,697 Train Loss: 0.0480297, Val Loss: 0.0654461
+2024-11-02 08:19:46,697 Epoch 186/500
+2024-11-02 08:23:20,862 Train Loss: 0.0480581, Val Loss: 0.0673987
+2024-11-02 08:23:20,862 Epoch 187/500
+2024-11-02 08:26:52,726 Train Loss: 0.0481855, Val Loss: 0.0666781
+2024-11-02 08:26:52,727 Epoch 188/500
+2024-11-02 08:29:48,905 Train Loss: 0.0484823, Val Loss: 0.0646197
+2024-11-02 08:29:48,906 Epoch 189/500
+2024-11-02 08:32:44,729 Train Loss: 0.0486292, Val Loss: 0.0667414
+2024-11-02 08:32:44,729 Epoch 190/500
+2024-11-02 08:36:05,764 Train Loss: 0.0486706, Val Loss: 0.0637663
+2024-11-02 08:36:05,764 Epoch 191/500
+2024-11-02 08:39:33,395 Train Loss: 0.0490878, Val Loss: 0.0633622
+2024-11-02 08:39:33,395 Epoch 192/500
+2024-11-02 08:42:28,960 Train Loss: 0.0478180, Val Loss: 0.0631506
+2024-11-02 08:42:28,961 Epoch 193/500
+2024-11-02 08:45:24,319 Train Loss: 0.0476836, Val Loss: 0.0635802
+2024-11-02 08:45:24,320 Epoch 194/500
+2024-11-02 08:48:53,799 Train Loss: 0.0476446, Val Loss: 0.0637492
+2024-11-02 08:48:53,800 Epoch 195/500
+2024-11-02 08:52:11,795 Train Loss: 0.0473838, Val Loss: 0.0624369
+2024-11-02 08:52:11,795 Epoch 196/500
+2024-11-02 08:55:07,417 Train Loss: 0.0469079, Val Loss: 0.0624948
+2024-11-02 08:55:07,418 Epoch 197/500
+2024-11-02 08:58:05,082 Train Loss: 0.0466404, Val Loss: 0.0633061
+2024-11-02 08:58:05,082 Epoch 198/500
+2024-11-02 09:01:33,582 Train Loss: 0.0468184, Val Loss: 0.0630486
+2024-11-02 09:01:33,583 Epoch 199/500
+2024-11-02 09:04:50,847 Train Loss: 0.0469254, Val Loss: 0.0629758
+2024-11-02 09:04:50,847 Epoch 200/500
+2024-11-02 09:07:47,306 Train Loss: 0.0467479, Val Loss: 0.0639387
+2024-11-02 09:07:47,307 Epoch 201/500
+2024-11-02 09:11:07,587 Train Loss: 0.0476059, Val Loss: 0.0643309
+2024-11-02 09:11:07,587 Epoch 202/500
+2024-11-02 09:14:32,691 Train Loss: 0.0470280, Val Loss: 0.0645014
+2024-11-02 09:14:32,691 Epoch 203/500
+2024-11-02 09:17:27,006 Train Loss: 0.0467279, Val Loss: 0.0640839
+2024-11-02 09:17:27,006 Epoch 204/500
+2024-11-02 09:20:32,603 Train Loss: 0.0464497, Val Loss: 0.0623920
+2024-11-02 09:20:32,603 Epoch 205/500
+2024-11-02 09:24:11,428 Train Loss: 0.0462260, Val Loss: 0.0627812
+2024-11-02 09:24:11,428 Epoch 206/500
+2024-11-02 09:27:06,928 Train Loss: 0.0460674, Val Loss: 0.0627784
+2024-11-02 09:27:06,929 Epoch 207/500
+2024-11-02 09:30:01,700 Train Loss: 0.0466174, Val Loss: 0.0642107
+2024-11-02 09:30:01,700 Epoch 208/500
+2024-11-02 09:33:39,722 Train Loss: 0.0464228, Val Loss: 0.0643455
+2024-11-02 09:33:39,722 Epoch 209/500
+2024-11-02 09:36:46,911 Train Loss: 0.0472782, Val Loss: 0.0638260
+2024-11-02 09:36:46,911 Epoch 210/500
+2024-11-02 09:39:41,318 Train Loss: 0.0475625, Val Loss: 0.0642031
+2024-11-02 09:39:41,319 Epoch 211/500
+2024-11-02 09:42:57,921 Train Loss: 0.0468253, Val Loss: 0.0650402
+2024-11-02 09:42:57,921 Epoch 212/500
+2024-11-02 09:46:26,197 Train Loss: 0.0470743, Val Loss: 0.0656636
+2024-11-02 09:46:26,198 Epoch 213/500
+2024-11-02 09:49:23,816 Train Loss: 0.0463564, Val Loss: 0.0664696
+2024-11-02 09:49:23,816 Epoch 214/500
+2024-11-02 09:53:09,256 Train Loss: 0.0467241, Val Loss: 0.0639718
+2024-11-02 09:53:09,256 Epoch 215/500
+2024-11-02 09:56:04,259 Train Loss: 0.0464902, Val Loss: 0.0639500
+2024-11-02 09:56:04,259 Epoch 216/500
+2024-11-02 09:58:57,620 Train Loss: 0.0462107, Val Loss: 0.0656044
+2024-11-02 09:58:57,620 Epoch 217/500
+2024-11-02 10:01:51,349 Train Loss: 0.0463994, Val Loss: 0.0677382
+2024-11-02 10:01:51,349 Epoch 218/500
+2024-11-02 10:05:00,887 Train Loss: 0.0466101, Val Loss: 0.0665351
+2024-11-02 10:05:00,887 Epoch 219/500
+2024-11-02 10:08:36,590 Train Loss: 0.0467272, Val Loss: 0.0654004
+2024-11-02 10:08:36,591 Epoch 220/500
+2024-11-02 10:11:33,107 Train Loss: 0.0470429, Val Loss: 0.0643636
+2024-11-02 10:11:33,107 Epoch 221/500
+2024-11-02 10:15:10,544 Train Loss: 0.0473663, Val Loss: 0.0636239
+2024-11-02 10:15:10,545 Epoch 222/500
+2024-11-02 10:18:19,683 Train Loss: 0.0468482, Val Loss: 0.0644520
+2024-11-02 10:18:19,683 Epoch 223/500
+2024-11-02 10:21:15,355 Train Loss: 0.0465695, Val Loss: 0.0641441
+2024-11-02 10:21:15,355 Epoch 224/500
+2024-11-02 10:24:57,099 Train Loss: 0.0459439, Val Loss: 0.0636620
+2024-11-02 10:24:57,099 Epoch 225/500
+2024-11-02 10:28:00,668 Train Loss: 0.0457703, Val Loss: 0.0644640
+2024-11-02 10:28:00,668 Epoch 226/500
+2024-11-02 10:30:57,621 Train Loss: 0.0452521, Val Loss: 0.0636580
+2024-11-02 10:30:57,621 Epoch 227/500
+2024-11-02 10:34:38,824 Train Loss: 0.0461391, Val Loss: 0.0628136
+2024-11-02 10:34:38,824 Epoch 228/500
+2024-11-02 10:37:41,907 Train Loss: 0.0445184, Val Loss: 0.0619271
+2024-11-02 10:37:41,907 Epoch 229/500
+2024-11-02 10:40:37,588 Train Loss: 0.0442432, Val Loss: 0.0627570
+2024-11-02 10:40:37,588 Epoch 230/500
+2024-11-02 10:43:35,210 Train Loss: 0.0439133, Val Loss: 0.0619078
+2024-11-02 10:43:35,210 Epoch 231/500
+2024-11-02 10:47:01,896 Train Loss: 0.0443042, Val Loss: 0.0618501
+2024-11-02 10:47:01,896 Epoch 232/500
+2024-11-02 10:50:22,728 Train Loss: 0.0445646, Val Loss: 0.0618263
+2024-11-02 10:50:22,728 Epoch 233/500
+2024-11-02 10:53:17,480 Train Loss: 0.0447499, Val Loss: 0.0619369
+2024-11-02 10:53:17,481 Epoch 234/500
+2024-11-02 10:56:12,094 Train Loss: 0.0444229, Val Loss: 0.0623105
+2024-11-02 10:56:12,094 Epoch 235/500
+2024-11-02 10:59:07,060 Train Loss: 0.0452141, Val Loss: 0.0624512
+2024-11-02 10:59:07,061 Epoch 236/500
+2024-11-02 11:02:02,684 Train Loss: 0.0451271, Val Loss: 0.0612406
+2024-11-02 11:02:02,685 Epoch 237/500
+2024-11-02 11:04:57,798 Train Loss: 0.0457152, Val Loss: 0.0630288
+2024-11-02 11:04:57,799 Epoch 238/500
+2024-11-02 11:07:51,540 Train Loss: 0.0458026, Val Loss: 0.0632026
+2024-11-02 11:07:51,541 Epoch 239/500
+2024-11-02 11:11:33,098 Train Loss: 0.0461623, Val Loss: 0.0626708
+2024-11-02 11:11:33,098 Epoch 240/500
+2024-11-02 11:14:37,824 Train Loss: 0.0462476, Val Loss: 0.0638611
+2024-11-02 11:14:37,824 Epoch 241/500
+2024-11-02 11:17:32,953 Train Loss: 0.0462913, Val Loss: 0.0634115
+2024-11-02 11:17:32,954 Epoch 242/500
+2024-11-02 11:20:34,928 Train Loss: 0.0463102, Val Loss: 0.0619711
+2024-11-02 11:20:34,929 Epoch 243/500
+2024-11-02 11:24:18,892 Train Loss: 0.0455850, Val Loss: 0.0623656
+2024-11-02 11:24:18,893 Epoch 244/500
+2024-11-02 11:27:15,059 Train Loss: 0.0450101, Val Loss: 0.0621282
+2024-11-02 11:27:15,059 Epoch 245/500
+2024-11-02 11:30:35,128 Train Loss: 0.0445481, Val Loss: 0.0608157
+2024-11-02 11:30:35,128 Epoch 246/500
+2024-11-02 11:34:02,769 Train Loss: 0.0439732, Val Loss: 0.0600980
+2024-11-02 11:34:02,769 Epoch 247/500
+2024-11-02 11:36:59,101 Train Loss: 0.0433323, Val Loss: 0.0602155
+2024-11-02 11:36:59,102 Epoch 248/500
+2024-11-02 11:39:55,311 Train Loss: 0.0438761, Val Loss: 0.0608776
+2024-11-02 11:39:55,311 Epoch 249/500
+2024-11-02 11:43:38,910 Train Loss: 0.0435855, Val Loss: 0.0602582
+2024-11-02 11:43:38,910 Epoch 250/500
+2024-11-02 11:46:44,047 Train Loss: 0.0436407, Val Loss: 0.0607103
+2024-11-02 11:46:44,047 Epoch 251/500
+2024-11-02 11:49:55,315 Train Loss: 0.0438314, Val Loss: 0.0602874
+2024-11-02 11:49:55,315 Epoch 252/500
+2024-11-02 11:53:29,607 Train Loss: 0.0434018, Val Loss: 0.0606785
+2024-11-02 11:53:29,607 Epoch 253/500
+2024-11-02 11:56:25,404 Train Loss: 0.0436848, Val Loss: 0.0602507
+2024-11-02 11:56:25,405 Epoch 254/500
+2024-11-02 11:59:53,906 Train Loss: 0.0432789, Val Loss: 0.0607832
+2024-11-02 11:59:53,907 Epoch 255/500
+2024-11-02 12:03:14,922 Train Loss: 0.0435869, Val Loss: 0.0613487
+2024-11-02 12:03:14,923 Epoch 256/500
+2024-11-02 12:06:12,148 Train Loss: 0.0440070, Val Loss: 0.0605806
+2024-11-02 12:06:12,148 Epoch 257/500
+2024-11-02 12:09:38,145 Train Loss: 0.0438834, Val Loss: 0.0606783
+2024-11-02 12:09:38,145 Epoch 258/500
+2024-11-02 12:13:05,703 Train Loss: 0.0432112, Val Loss: 0.0607154
+2024-11-02 12:13:05,703 Epoch 259/500
+2024-11-02 12:16:04,319 Train Loss: 0.0435320, Val Loss: 0.0602177
+2024-11-02 12:16:04,319 Epoch 260/500
+2024-11-02 12:19:44,577 Train Loss: 0.0442541, Val Loss: 0.0605268
+2024-11-02 12:19:44,578 Epoch 261/500
+2024-11-02 12:23:06,594 Train Loss: 0.0436440, Val Loss: 0.0608823
+2024-11-02 12:23:06,594 Epoch 262/500
+2024-11-02 12:26:40,181 Train Loss: 0.0439424, Val Loss: 0.0609285
+2024-11-02 12:26:40,182 Epoch 263/500
+2024-11-02 12:30:07,036 Train Loss: 0.0441003, Val Loss: 0.0614702
+2024-11-02 12:30:07,037 Epoch 264/500
+2024-11-02 12:33:17,103 Train Loss: 0.0443392, Val Loss: 0.0613654
+2024-11-02 12:33:17,104 Epoch 265/500
+2024-11-02 12:36:58,419 Train Loss: 0.0449716, Val Loss: 0.0604724
+2024-11-02 12:36:58,420 Epoch 266/500
+2024-11-02 12:39:58,737 Train Loss: 0.0452770, Val Loss: 0.0606473
+2024-11-02 12:39:58,737 Epoch 267/500
+2024-11-02 12:43:20,531 Train Loss: 0.0445917, Val Loss: 0.0604431
+2024-11-02 12:43:20,532 Epoch 268/500
+2024-11-02 12:46:53,001 Train Loss: 0.0445921, Val Loss: 0.0606213
+2024-11-02 12:46:53,001 Epoch 269/500
+2024-11-02 12:49:52,948 Train Loss: 0.0442547, Val Loss: 0.0610656
+2024-11-02 12:49:52,948 Epoch 270/500
+2024-11-02 12:53:40,502 Train Loss: 0.0436127, Val Loss: 0.0616169
+2024-11-02 12:53:40,503 Epoch 271/500
+2024-11-02 12:56:45,999 Train Loss: 0.0437103, Val Loss: 0.0608491
+2024-11-02 12:56:45,999 Epoch 272/500
+2024-11-02 12:59:45,955 Train Loss: 0.0427906, Val Loss: 0.0613495
+2024-11-02 12:59:45,956 Epoch 273/500
+2024-11-02 13:03:29,888 Train Loss: 0.0426337, Val Loss: 0.0601556
+2024-11-02 13:03:29,889 Epoch 274/500
+2024-11-02 13:06:41,654 Train Loss: 0.0434502, Val Loss: 0.0593469
+2024-11-02 13:06:41,654 Epoch 275/500
+2024-11-02 13:09:42,127 Train Loss: 0.0434018, Val Loss: 0.0595080
+2024-11-02 13:09:42,128 Epoch 276/500
+2024-11-02 13:12:47,399 Train Loss: 0.0432626, Val Loss: 0.0602453
+2024-11-02 13:12:47,399 Epoch 277/500
+2024-11-02 13:16:23,394 Train Loss: 0.0431200, Val Loss: 0.0601694
+2024-11-02 13:16:23,395 Epoch 278/500
+2024-11-02 13:19:50,280 Train Loss: 0.0423720, Val Loss: 0.0600633
+2024-11-02 13:19:50,281 Epoch 279/500
+2024-11-02 13:23:05,036 Train Loss: 0.0419733, Val Loss: 0.0589951
+2024-11-02 13:23:05,037 Epoch 280/500
+2024-11-02 13:26:44,725 Train Loss: 0.0420052, Val Loss: 0.0590216
+2024-11-02 13:26:44,726 Epoch 281/500
+2024-11-02 13:29:44,214 Train Loss: 0.0416629, Val Loss: 0.0587665
+2024-11-02 13:29:44,214 Epoch 282/500
+2024-11-02 13:32:42,448 Train Loss: 0.0415587, Val Loss: 0.0591061
+2024-11-02 13:32:42,449 Epoch 283/500
+2024-11-02 13:36:37,899 Train Loss: 0.0417753, Val Loss: 0.0589201
+2024-11-02 13:36:37,899 Epoch 284/500
+2024-11-02 13:39:39,564 Train Loss: 0.0417630, Val Loss: 0.0590874
+2024-11-02 13:39:39,564 Epoch 285/500
+2024-11-02 13:42:39,028 Train Loss: 0.0415218, Val Loss: 0.0591461
+2024-11-02 13:42:39,028 Epoch 286/500
+2024-11-02 13:45:42,033 Train Loss: 0.0417465, Val Loss: 0.0600360
+2024-11-02 13:45:42,033 Epoch 287/500
+2024-11-02 13:49:32,992 Train Loss: 0.0420417, Val Loss: 0.0601999
+2024-11-02 13:49:32,992 Epoch 288/500
+2024-11-02 13:52:33,846 Train Loss: 0.0420656, Val Loss: 0.0598516
+2024-11-02 13:52:33,846 Epoch 289/500
+2024-11-02 13:55:32,240 Train Loss: 0.0415406, Val Loss: 0.0591464
+2024-11-02 13:55:32,240 Epoch 290/500
+2024-11-02 13:58:33,550 Train Loss: 0.0425317, Val Loss: 0.0601851
+2024-11-02 13:58:33,551 Epoch 291/500
+2024-11-02 14:01:55,772 Train Loss: 0.0421410, Val Loss: 0.0604017
+2024-11-02 14:01:55,772 Epoch 292/500
+2024-11-02 14:05:30,545 Train Loss: 0.0417393, Val Loss: 0.0591543
+2024-11-02 14:05:30,545 Epoch 293/500
+2024-11-02 14:08:32,246 Train Loss: 0.0415275, Val Loss: 0.0585560
+2024-11-02 14:08:32,246 Epoch 294/500
+2024-11-02 14:11:31,316 Train Loss: 0.0414154, Val Loss: 0.0591052
+2024-11-02 14:11:31,316 Epoch 295/500
+2024-11-02 14:14:30,206 Train Loss: 0.0420112, Val Loss: 0.0602253
+2024-11-02 14:14:30,206 Epoch 296/500
+2024-11-02 14:17:28,202 Train Loss: 0.0422841, Val Loss: 0.0610316
+2024-11-02 14:17:28,202 Epoch 297/500
+2024-11-02 14:20:26,173 Train Loss: 0.0428998, Val Loss: 0.0588401
+2024-11-02 14:20:26,173 Epoch 298/500
+2024-11-02 14:23:47,174 Train Loss: 0.0422060, Val Loss: 0.0594055
+2024-11-02 14:23:47,174 Epoch 299/500
+2024-11-02 14:27:20,952 Train Loss: 0.0416644, Val Loss: 0.0595470
+2024-11-02 14:27:20,953 Epoch 300/500
+2024-11-02 14:30:19,386 Train Loss: 0.0418225, Val Loss: 0.0589879
+2024-11-02 14:30:19,387 Epoch 301/500
+2024-11-02 14:33:15,452 Train Loss: 0.0415816, Val Loss: 0.0594943
+2024-11-02 14:33:15,453 Epoch 302/500
+2024-11-02 14:36:50,207 Train Loss: 0.0419323, Val Loss: 0.0594391
+2024-11-02 14:36:50,207 Epoch 303/500
+2024-11-02 14:40:07,813 Train Loss: 0.0419379, Val Loss: 0.0600717
+2024-11-02 14:40:07,814 Epoch 304/500
+2024-11-02 14:43:13,206 Train Loss: 0.0423009, Val Loss: 0.0595942
+2024-11-02 14:43:13,207 Epoch 305/500
+2024-11-02 14:46:16,218 Train Loss: 0.0416932, Val Loss: 0.0597923
+2024-11-02 14:46:16,218 Epoch 306/500
+2024-11-02 14:49:19,229 Train Loss: 0.0418825, Val Loss: 0.0597237
+2024-11-02 14:49:19,229 Epoch 307/500
+2024-11-02 14:52:23,987 Train Loss: 0.0419184, Val Loss: 0.0602407
+2024-11-02 14:52:23,988 Epoch 308/500
+2024-11-02 14:55:31,509 Train Loss: 0.0419676, Val Loss: 0.0610514
+2024-11-02 14:55:31,509 Epoch 309/500
+2024-11-02 14:59:19,151 Train Loss: 0.0422856, Val Loss: 0.0615857
+2024-11-02 14:59:19,152 Epoch 310/500
+2024-11-02 15:02:17,619 Train Loss: 0.0419358, Val Loss: 0.0618955
+2024-11-02 15:02:17,619 Epoch 311/500
+2024-11-02 15:05:15,218 Train Loss: 0.0419437, Val Loss: 0.0604494
+2024-11-02 15:05:15,219 Epoch 312/500
+2024-11-02 15:08:13,784 Train Loss: 0.0415363, Val Loss: 0.0605207
+2024-11-02 15:08:13,785 Epoch 313/500
+2024-11-02 15:11:13,682 Train Loss: 0.0415161, Val Loss: 0.0600125
+2024-11-02 15:11:13,682 Epoch 314/500
+2024-11-02 15:14:12,237 Train Loss: 0.0410120, Val Loss: 0.0596264
+2024-11-02 15:14:12,238 Epoch 315/500
+2024-11-02 15:17:24,820 Train Loss: 0.0412453, Val Loss: 0.0588437
+2024-11-02 15:17:24,820 Epoch 316/500
+2024-11-02 15:21:06,222 Train Loss: 0.0407807, Val Loss: 0.0597100
+2024-11-02 15:21:06,222 Epoch 317/500
+2024-11-02 15:24:04,940 Train Loss: 0.0406798, Val Loss: 0.0594845
+2024-11-02 15:24:04,940 Epoch 318/500
+2024-11-02 15:27:03,695 Train Loss: 0.0409038, Val Loss: 0.0593843
+2024-11-02 15:27:03,695 Epoch 319/500
+2024-11-02 15:30:01,087 Train Loss: 0.0403016, Val Loss: 0.0592302
+2024-11-02 15:30:01,087 Epoch 320/500
+2024-11-02 15:33:03,522 Train Loss: 0.0410438, Val Loss: 0.0602990
+2024-11-02 15:33:03,522 Epoch 321/500
+2024-11-02 15:36:01,446 Train Loss: 0.0409603, Val Loss: 0.0601965
+2024-11-02 15:36:01,446 Epoch 322/500
+2024-11-02 15:39:01,771 Train Loss: 0.0412127, Val Loss: 0.0597409
+2024-11-02 15:39:01,771 Epoch 323/500
+2024-11-02 15:42:01,298 Train Loss: 0.0417300, Val Loss: 0.0597532
+2024-11-02 15:42:01,299 Epoch 324/500
+2024-11-02 15:45:30,788 Train Loss: 0.0418059, Val Loss: 0.0599914
+2024-11-02 15:45:30,788 Epoch 325/500
+2024-11-02 15:48:56,593 Train Loss: 0.0418212, Val Loss: 0.0594606
+2024-11-02 15:48:56,593 Epoch 326/500
+2024-11-02 15:51:55,467 Train Loss: 0.0411545, Val Loss: 0.0596464
+2024-11-02 15:51:55,468 Epoch 327/500
+2024-11-02 15:55:28,297 Train Loss: 0.0410273, Val Loss: 0.0599024
+2024-11-02 15:55:28,298 Epoch 328/500
+2024-11-02 15:58:49,204 Train Loss: 0.0407685, Val Loss: 0.0587076
+2024-11-02 15:58:49,204 Epoch 329/500
+2024-11-02 16:01:56,308 Train Loss: 0.0405287, Val Loss: 0.0585450
+2024-11-02 16:01:56,308 Epoch 330/500
+2024-11-02 16:05:38,900 Train Loss: 0.0408748, Val Loss: 0.0579685
+2024-11-02 16:05:38,900 Epoch 331/500
+2024-11-02 16:08:39,207 Train Loss: 0.0409528, Val Loss: 0.0580340
+2024-11-02 16:08:39,208 Epoch 332/500
+2024-11-02 16:11:38,415 Train Loss: 0.0413668, Val Loss: 0.0586145
+2024-11-02 16:11:38,416 Epoch 333/500
+2024-11-02 16:15:30,627 Train Loss: 0.0415958, Val Loss: 0.0582683
+2024-11-02 16:15:30,628 Epoch 334/500
+2024-11-02 16:18:28,255 Train Loss: 0.0419676, Val Loss: 0.0592931
+2024-11-02 16:18:28,256 Epoch 335/500
+2024-11-02 16:21:24,581 Train Loss: 0.0414313, Val Loss: 0.0581618
+2024-11-02 16:21:24,581 Epoch 336/500
+2024-11-02 16:24:24,059 Train Loss: 0.0414570, Val Loss: 0.0577461
+2024-11-02 16:24:24,059 Epoch 337/500
+2024-11-02 16:27:22,446 Train Loss: 0.0423431, Val Loss: 0.0578982
+2024-11-02 16:27:22,447 Epoch 338/500
+2024-11-02 16:30:19,873 Train Loss: 0.0419595, Val Loss: 0.0576586
+2024-11-02 16:30:19,873 Epoch 339/500
+2024-11-02 16:33:56,030 Train Loss: 0.0411537, Val Loss: 0.0575664
+2024-11-02 16:33:56,030 Epoch 340/500
+2024-11-02 16:37:12,219 Train Loss: 0.0410217, Val Loss: 0.0573184
+2024-11-02 16:37:12,219 Epoch 341/500
+2024-11-02 16:40:13,544 Train Loss: 0.0410328, Val Loss: 0.0578781
+2024-11-02 16:40:13,545 Epoch 342/500
+2024-11-02 16:43:18,809 Train Loss: 0.0411319, Val Loss: 0.0586561
+2024-11-02 16:43:18,809 Epoch 343/500
+2024-11-02 16:47:08,796 Train Loss: 0.0411878, Val Loss: 0.0590822
+2024-11-02 16:47:08,796 Epoch 344/500
+2024-11-02 16:50:10,317 Train Loss: 0.0409471, Val Loss: 0.0590031
+2024-11-02 16:50:10,318 Epoch 345/500
+2024-11-02 16:53:07,772 Train Loss: 0.0403242, Val Loss: 0.0582457
+2024-11-02 16:53:07,772 Epoch 346/500
+2024-11-02 16:56:06,252 Train Loss: 0.0398690, Val Loss: 0.0583578
+2024-11-02 16:56:06,252 Epoch 347/500
+2024-11-02 16:59:22,260 Train Loss: 0.0397748, Val Loss: 0.0579849
+2024-11-02 16:59:22,260 Epoch 348/500
+2024-11-02 17:03:01,150 Train Loss: 0.0400483, Val Loss: 0.0584467
+2024-11-02 17:03:01,151 Epoch 349/500
+2024-11-02 17:05:59,962 Train Loss: 0.0401085, Val Loss: 0.0580592
+2024-11-02 17:05:59,962 Epoch 350/500
+2024-11-02 17:09:43,899 Train Loss: 0.0397280, Val Loss: 0.0584481
+2024-11-02 17:09:43,900 Epoch 351/500
+2024-11-02 17:12:52,647 Train Loss: 0.0403896, Val Loss: 0.0580628
+2024-11-02 17:12:52,648 Epoch 352/500
+2024-11-02 17:15:53,901 Train Loss: 0.0398922, Val Loss: 0.0579900
+2024-11-02 17:15:53,902 Epoch 353/500
+2024-11-02 17:19:40,538 Train Loss: 0.0399874, Val Loss: 0.0584608
+2024-11-02 17:19:40,539 Epoch 354/500
+2024-11-02 17:22:44,212 Train Loss: 0.0401203, Val Loss: 0.0573071
+2024-11-02 17:22:44,212 Epoch 355/500
+2024-11-02 17:25:42,570 Train Loss: 0.0406190, Val Loss: 0.0592341
+2024-11-02 17:25:42,571 Epoch 356/500
+2024-11-02 17:28:41,767 Train Loss: 0.0408931, Val Loss: 0.0589838
+2024-11-02 17:28:41,768 Epoch 357/500
+2024-11-02 17:32:07,679 Train Loss: 0.0409831, Val Loss: 0.0590234
+2024-11-02 17:32:07,679 Epoch 358/500
+2024-11-02 17:35:29,489 Train Loss: 0.0405867, Val Loss: 0.0579176
+2024-11-02 17:35:29,489 Epoch 359/500
+2024-11-02 17:38:27,832 Train Loss: 0.0405270, Val Loss: 0.0578312
+2024-11-02 17:38:27,834 Epoch 360/500
+2024-11-02 17:42:03,325 Train Loss: 0.0410094, Val Loss: 0.0579466
+2024-11-02 17:42:03,325 Epoch 361/500
+2024-11-02 17:45:16,224 Train Loss: 0.0408926, Val Loss: 0.0582639
+2024-11-02 17:45:16,224 Epoch 362/500
+2024-11-02 17:48:35,811 Train Loss: 0.0409729, Val Loss: 0.0580129
+2024-11-02 17:48:35,811 Epoch 363/500
+2024-11-02 17:52:07,111 Train Loss: 0.0404854, Val Loss: 0.0581374
+2024-11-02 17:52:07,111 Epoch 364/500
+2024-11-02 17:55:04,740 Train Loss: 0.0402474, Val Loss: 0.0577086
+2024-11-02 17:55:04,740 Epoch 365/500
+2024-11-02 17:58:31,194 Train Loss: 0.0400140, Val Loss: 0.0573770
+2024-11-02 17:58:31,194 Epoch 366/500
+2024-11-02 18:01:54,779 Train Loss: 0.0394420, Val Loss: 0.0575861
+2024-11-02 18:01:54,779 Epoch 367/500
+2024-11-02 18:04:51,633 Train Loss: 0.0390195, Val Loss: 0.0571440
+2024-11-02 18:04:51,633 Epoch 368/500
+2024-11-02 18:07:48,145 Train Loss: 0.0390252, Val Loss: 0.0575364
+2024-11-02 18:07:48,145 Epoch 369/500
+2024-11-02 18:10:45,667 Train Loss: 0.0391279, Val Loss: 0.0573542
+2024-11-02 18:10:45,668 Epoch 370/500
+2024-11-02 18:14:32,875 Train Loss: 0.0392613, Val Loss: 0.0577171
+2024-11-02 18:14:32,876 Epoch 371/500
+2024-11-02 18:18:05,367 Train Loss: 0.0390012, Val Loss: 0.0582181
+2024-11-02 18:18:05,368 Epoch 372/500
+2024-11-02 18:21:28,554 Train Loss: 0.0396738, Val Loss: 0.0588587
+2024-11-02 18:21:28,554 Epoch 373/500
+2024-11-02 18:24:25,270 Train Loss: 0.0394583, Val Loss: 0.0588800
+2024-11-02 18:24:25,271 Epoch 374/500
+2024-11-02 18:27:40,406 Train Loss: 0.0399388, Val Loss: 0.0589896
+2024-11-02 18:27:40,406 Epoch 375/500
+2024-11-02 18:31:16,299 Train Loss: 0.0399578, Val Loss: 0.0598740
+2024-11-02 18:31:16,300 Epoch 376/500
+2024-11-02 18:34:16,167 Train Loss: 0.0394974, Val Loss: 0.0580215
+2024-11-02 18:34:16,167 Epoch 377/500
+2024-11-02 18:38:00,190 Train Loss: 0.0396564, Val Loss: 0.0579727
+2024-11-02 18:38:00,191 Epoch 378/500
+2024-11-02 18:41:11,688 Train Loss: 0.0402715, Val Loss: 0.0584190
+2024-11-02 18:41:11,689 Epoch 379/500
+2024-11-02 18:44:11,388 Train Loss: 0.0405642, Val Loss: 0.0580653
+2024-11-02 18:44:11,389 Epoch 380/500
+2024-11-02 18:47:11,544 Train Loss: 0.0398939, Val Loss: 0.0587236
+2024-11-02 18:47:11,545 Epoch 381/500
+2024-11-02 18:50:46,785 Train Loss: 0.0401086, Val Loss: 0.0597313
+2024-11-02 18:50:46,786 Epoch 382/500
+2024-11-02 18:54:06,428 Train Loss: 0.0401842, Val Loss: 0.0590873
+2024-11-02 18:54:06,429 Epoch 383/500
+2024-11-02 18:57:07,851 Train Loss: 0.0403904, Val Loss: 0.0588594
+2024-11-02 18:57:07,851 Epoch 384/500
+2024-11-02 19:00:57,965 Train Loss: 0.0406806, Val Loss: 0.0588065
+2024-11-02 19:00:57,965 Epoch 385/500
+2024-11-02 19:04:06,347 Train Loss: 0.0405931, Val Loss: 0.0579386
+2024-11-02 19:04:06,347 Epoch 386/500
+2024-11-02 19:07:15,003 Train Loss: 0.0406293, Val Loss: 0.0585173
+2024-11-02 19:07:15,003 Epoch 387/500
+2024-11-02 19:10:16,047 Train Loss: 0.0397994, Val Loss: 0.0583322
+2024-11-02 19:10:16,048 Epoch 388/500
+2024-11-02 19:13:42,369 Train Loss: 0.0401114, Val Loss: 0.0588679
+2024-11-02 19:13:42,369 Epoch 389/500
+2024-11-02 19:17:06,873 Train Loss: 0.0400305, Val Loss: 0.0579988
+2024-11-02 19:17:06,873 Epoch 390/500
+2024-11-02 19:20:03,609 Train Loss: 0.0400819, Val Loss: 0.0579372
+2024-11-02 19:20:03,610 Epoch 391/500
+2024-11-02 19:23:40,201 Train Loss: 0.0400527, Val Loss: 0.0584015
+2024-11-02 19:23:40,202 Epoch 392/500
+2024-11-02 19:27:14,838 Train Loss: 0.0399731, Val Loss: 0.0586223
+2024-11-02 19:27:14,838 Epoch 393/500
+2024-11-02 19:30:13,107 Train Loss: 0.0393839, Val Loss: 0.0579094
+2024-11-02 19:30:13,108 Epoch 394/500
+2024-11-02 19:33:10,386 Train Loss: 0.0397058, Val Loss: 0.0577166
+2024-11-02 19:33:10,387 Epoch 395/500
+2024-11-02 19:36:53,632 Train Loss: 0.0389187, Val Loss: 0.0574457
+2024-11-02 19:36:53,632 Epoch 396/500
+2024-11-02 19:39:58,259 Train Loss: 0.0393241, Val Loss: 0.0581951
+2024-11-02 19:39:58,259 Epoch 397/500
+2024-11-02 19:42:54,344 Train Loss: 0.0387192, Val Loss: 0.0579340
+2024-11-02 19:42:54,344 Epoch 398/500
+2024-11-02 19:46:24,509 Train Loss: 0.0388178, Val Loss: 0.0577731
+2024-11-02 19:46:24,509 Epoch 399/500
+2024-11-02 19:49:46,061 Train Loss: 0.0389307, Val Loss: 0.0578627
+2024-11-02 19:49:46,062 Epoch 400/500
+2024-11-02 19:52:44,098 Train Loss: 0.0386144, Val Loss: 0.0579973
+2024-11-02 19:52:44,098 Epoch 401/500
+2024-11-02 19:56:05,667 Train Loss: 0.0392295, Val Loss: 0.0579757
+2024-11-02 19:56:05,667 Epoch 402/500
+2024-11-02 19:59:32,372 Train Loss: 0.0388906, Val Loss: 0.0591460
+2024-11-02 19:59:32,373 Epoch 403/500
+2024-11-02 20:02:30,281 Train Loss: 0.0383377, Val Loss: 0.0599615
+2024-11-02 20:02:30,281 Epoch 404/500
+2024-11-02 20:06:06,921 Train Loss: 0.0388708, Val Loss: 0.0586047
+2024-11-02 20:06:06,921 Epoch 405/500
+2024-11-02 20:09:21,189 Train Loss: 0.0387721, Val Loss: 0.0578503
+2024-11-02 20:09:21,190 Epoch 406/500
+2024-11-02 20:12:18,575 Train Loss: 0.0389045, Val Loss: 0.0578896
+2024-11-02 20:12:18,576 Epoch 407/500
+2024-11-02 20:15:30,692 Train Loss: 0.0390265, Val Loss: 0.0577000
+2024-11-02 20:15:30,692 Epoch 408/500
+2024-11-02 20:19:08,119 Train Loss: 0.0393684, Val Loss: 0.0574350
+2024-11-02 20:19:08,119 Epoch 409/500
+2024-11-02 20:22:09,171 Train Loss: 0.0395314, Val Loss: 0.0575468
+2024-11-02 20:22:09,172 Epoch 410/500
+2024-11-02 20:25:39,665 Train Loss: 0.0391024, Val Loss: 0.0580800
+2024-11-02 20:25:39,666 Epoch 411/500
+2024-11-02 20:29:03,500 Train Loss: 0.0390637, Val Loss: 0.0579825
+2024-11-02 20:29:03,500 Epoch 412/500
+2024-11-02 20:32:10,281 Train Loss: 0.0389488, Val Loss: 0.0583745
+2024-11-02 20:32:10,282 Epoch 413/500
+2024-11-02 20:35:48,152 Train Loss: 0.0388941, Val Loss: 0.0577608
+2024-11-02 20:35:48,152 Epoch 414/500
+2024-11-02 20:39:07,289 Train Loss: 0.0388513, Val Loss: 0.0578232
+2024-11-02 20:39:07,289 Epoch 415/500
+2024-11-02 20:42:15,604 Train Loss: 0.0390740, Val Loss: 0.0584172
+2024-11-02 20:42:15,604 Epoch 416/500
+2024-11-02 20:45:59,260 Train Loss: 0.0393258, Val Loss: 0.0587398
+2024-11-02 20:45:59,261 Epoch 417/500
+2024-11-02 20:48:59,119 Train Loss: 0.0397258, Val Loss: 0.0582443
+2024-11-02 20:48:59,119 Epoch 418/500
+2024-11-02 20:52:20,613 Train Loss: 0.0402310, Val Loss: 0.0576836
+2024-11-02 20:52:20,613 Epoch 419/500
+2024-11-02 20:55:48,145 Train Loss: 0.0392868, Val Loss: 0.0577936
+2024-11-02 20:55:48,146 Epoch 420/500
+2024-11-02 20:58:46,201 Train Loss: 0.0398394, Val Loss: 0.0568674
+2024-11-02 20:58:46,201 Epoch 421/500
+2024-11-02 21:01:45,445 Train Loss: 0.0391285, Val Loss: 0.0567561
+2024-11-02 21:01:45,445 Epoch 422/500
+2024-11-02 21:04:42,728 Train Loss: 0.0391262, Val Loss: 0.0567382
+2024-11-02 21:04:42,728 Epoch 423/500
+2024-11-02 21:08:27,228 Train Loss: 0.0387613, Val Loss: 0.0571007
+2024-11-02 21:08:27,228 Epoch 424/500
+2024-11-02 21:11:35,389 Train Loss: 0.0388678, Val Loss: 0.0569020
+2024-11-02 21:11:35,389 Epoch 425/500
+2024-11-02 21:14:49,880 Train Loss: 0.0387865, Val Loss: 0.0574812
+2024-11-02 21:14:49,881 Epoch 426/500
+2024-11-02 21:18:25,639 Train Loss: 0.0387229, Val Loss: 0.0571121
+2024-11-02 21:18:25,640 Epoch 427/500
+2024-11-02 21:21:45,145 Train Loss: 0.0390190, Val Loss: 0.0567339
+2024-11-02 21:21:45,146 Epoch 428/500
+2024-11-02 21:25:23,734 Train Loss: 0.0392649, Val Loss: 0.0566259
+2024-11-02 21:25:23,734 Epoch 429/500
+2024-11-02 21:28:21,801 Train Loss: 0.0388042, Val Loss: 0.0570679
+2024-11-02 21:28:21,802 Epoch 430/500
+2024-11-02 21:31:20,899 Train Loss: 0.0388690, Val Loss: 0.0575413
+2024-11-02 21:31:20,900 Epoch 431/500
+2024-11-02 21:35:06,997 Train Loss: 0.0383034, Val Loss: 0.0579443
+2024-11-02 21:35:06,998 Epoch 432/500
+2024-11-02 21:38:14,242 Train Loss: 0.0381320, Val Loss: 0.0577746
+2024-11-02 21:38:14,242 Epoch 433/500
+2024-11-02 21:41:17,206 Train Loss: 0.0388885, Val Loss: 0.0582790
+2024-11-02 21:41:17,206 Epoch 434/500
+2024-11-02 21:45:03,712 Train Loss: 0.0389449, Val Loss: 0.0576361
+2024-11-02 21:45:03,712 Epoch 435/500
+2024-11-02 21:48:24,850 Train Loss: 0.0389969, Val Loss: 0.0576007
+2024-11-02 21:48:24,850 Epoch 436/500
+2024-11-02 21:51:53,680 Train Loss: 0.0391282, Val Loss: 0.0565848
+2024-11-02 21:51:53,681 Epoch 437/500
+2024-11-02 21:54:53,127 Train Loss: 0.0385026, Val Loss: 0.0564438
+2024-11-02 21:54:53,127 Epoch 438/500
+2024-11-02 21:57:54,417 Train Loss: 0.0386925, Val Loss: 0.0559600
+2024-11-02 21:57:54,417 Epoch 439/500
+2024-11-02 22:01:20,923 Train Loss: 0.0391002, Val Loss: 0.0560886
+2024-11-02 22:01:20,923 Epoch 440/500
+2024-11-02 22:04:46,216 Train Loss: 0.0395528, Val Loss: 0.0559410
+2024-11-02 22:04:46,216 Epoch 441/500
+2024-11-02 22:07:46,139 Train Loss: 0.0399341, Val Loss: 0.0560736
+2024-11-02 22:07:46,139 Epoch 442/500
+2024-11-02 22:10:44,509 Train Loss: 0.0398281, Val Loss: 0.0564665
+2024-11-02 22:10:44,510 Epoch 443/500
+2024-11-02 22:14:10,688 Train Loss: 0.0400747, Val Loss: 0.0560445
+2024-11-02 22:14:10,688 Epoch 444/500
+2024-11-02 22:17:34,967 Train Loss: 0.0409284, Val Loss: 0.0565227
+2024-11-02 22:17:34,968 Epoch 445/500
+2024-11-02 22:20:33,352 Train Loss: 0.0396965, Val Loss: 0.0570916
+2024-11-02 22:20:33,352 Epoch 446/500
+2024-11-02 22:24:20,103 Train Loss: 0.0395606, Val Loss: 0.0573105
+2024-11-02 22:24:20,104 Epoch 447/500
+2024-11-02 22:27:20,865 Train Loss: 0.0394130, Val Loss: 0.0568101
+2024-11-02 22:27:20,866 Epoch 448/500
+2024-11-02 22:30:16,627 Train Loss: 0.0387025, Val Loss: 0.0570552
+2024-11-02 22:30:16,627 Epoch 449/500
+2024-11-02 22:34:04,620 Train Loss: 0.0392206, Val Loss: 0.0567237
+2024-11-02 22:34:04,620 Epoch 450/500
+2024-11-02 22:37:06,377 Train Loss: 0.0384510, Val Loss: 0.0560127
+2024-11-02 22:37:06,377 Epoch 451/500
+2024-11-02 22:40:12,127 Train Loss: 0.0385708, Val Loss: 0.0567452
+2024-11-02 22:40:12,132 Epoch 452/500
+2024-11-02 22:43:56,311 Train Loss: 0.0382739, Val Loss: 0.0563417
+2024-11-02 22:43:56,312 Epoch 453/500
+2024-11-02 22:47:08,413 Train Loss: 0.0376996, Val Loss: 0.0558982
+2024-11-02 22:47:08,414 Epoch 454/500
+2024-11-02 22:50:49,669 Train Loss: 0.0374447, Val Loss: 0.0556690
+2024-11-02 22:50:49,670 Epoch 455/500
+2024-11-02 22:53:49,007 Train Loss: 0.0376442, Val Loss: 0.0558452
+2024-11-02 22:53:49,008 Epoch 456/500
+2024-11-02 22:56:47,089 Train Loss: 0.0376348, Val Loss: 0.0559104
+2024-11-02 22:56:47,090 Epoch 457/500
+2024-11-02 23:00:23,064 Train Loss: 0.0373049, Val Loss: 0.0559752
+2024-11-02 23:00:23,064 Epoch 458/500
+2024-11-02 23:03:36,628 Train Loss: 0.0373410, Val Loss: 0.0557501
+2024-11-02 23:03:36,628 Epoch 459/500
+2024-11-02 23:06:56,943 Train Loss: 0.0376213, Val Loss: 0.0554219
+2024-11-02 23:06:56,943 Epoch 460/500
+2024-11-02 23:10:25,010 Train Loss: 0.0374262, Val Loss: 0.0553996
+2024-11-02 23:10:25,010 Epoch 461/500
+2024-11-02 23:13:21,931 Train Loss: 0.0377059, Val Loss: 0.0553632
+2024-11-02 23:13:21,931 Epoch 462/500
+2024-11-02 23:16:19,013 Train Loss: 0.0369894, Val Loss: 0.0554736
+2024-11-02 23:16:19,013 Epoch 463/500
+2024-11-02 23:19:53,004 Train Loss: 0.0374426, Val Loss: 0.0555118
+2024-11-02 23:19:53,004 Epoch 464/500
+2024-11-02 23:23:10,115 Train Loss: 0.0374760, Val Loss: 0.0559440
+2024-11-02 23:23:10,115 Epoch 465/500
+2024-11-02 23:26:29,706 Train Loss: 0.0377081, Val Loss: 0.0563126
+2024-11-02 23:26:29,707 Epoch 466/500
+2024-11-02 23:29:59,147 Train Loss: 0.0377832, Val Loss: 0.0564361
+2024-11-02 23:29:59,147 Epoch 467/500
+2024-11-02 23:33:00,675 Train Loss: 0.0378480, Val Loss: 0.0565159
+2024-11-02 23:33:00,676 Epoch 468/500
+2024-11-02 23:36:47,620 Train Loss: 0.0381216, Val Loss: 0.0563773
+2024-11-02 23:36:47,620 Epoch 469/500
+2024-11-02 23:39:46,182 Train Loss: 0.0386798, Val Loss: 0.0571887
+2024-11-02 23:39:46,183 Epoch 470/500
+2024-11-02 23:42:44,607 Train Loss: 0.0386984, Val Loss: 0.0576652
+2024-11-02 23:42:44,607 Epoch 471/500
+2024-11-02 23:45:43,140 Train Loss: 0.0386614, Val Loss: 0.0570068
+2024-11-02 23:45:43,140 Epoch 472/500
+2024-11-02 23:48:59,929 Train Loss: 0.0386925, Val Loss: 0.0559145
+2024-11-02 23:48:59,929 Epoch 473/500
+2024-11-02 23:52:37,953 Train Loss: 0.0388576, Val Loss: 0.0563857
+2024-11-02 23:52:37,953 Epoch 474/500
+2024-11-02 23:55:50,616 Train Loss: 0.0383907, Val Loss: 0.0573166
+2024-11-02 23:55:50,617 Epoch 475/500
+2024-11-02 23:59:35,482 Train Loss: 0.0382544, Val Loss: 0.0566372
+2024-11-02 23:59:35,482 Epoch 476/500
+2024-11-03 00:03:13,219 Train Loss: 0.0385290, Val Loss: 0.0559754
+2024-11-03 00:03:13,219 Epoch 477/500
+2024-11-03 00:06:28,398 Train Loss: 0.0379065, Val Loss: 0.0557118
+2024-11-03 00:06:28,399 Epoch 478/500
+2024-11-03 00:09:26,315 Train Loss: 0.0383753, Val Loss: 0.0557908
+2024-11-03 00:09:26,316 Epoch 479/500
+2024-11-03 00:13:15,129 Train Loss: 0.0377758, Val Loss: 0.0557593
+2024-11-03 00:13:15,130 Epoch 480/500
+2024-11-03 00:16:12,726 Train Loss: 0.0373506, Val Loss: 0.0559643
+2024-11-03 00:16:12,727 Epoch 481/500
+2024-11-03 00:19:10,168 Train Loss: 0.0372199, Val Loss: 0.0566857
+2024-11-03 00:19:10,168 Epoch 482/500
+2024-11-03 00:22:07,301 Train Loss: 0.0371610, Val Loss: 0.0558072
+2024-11-03 00:22:07,301 Epoch 483/500
+2024-11-03 00:25:16,442 Train Loss: 0.0369303, Val Loss: 0.0557542
+2024-11-03 00:25:16,443 Epoch 484/500
+2024-11-03 00:28:55,691 Train Loss: 0.0367464, Val Loss: 0.0557234
+2024-11-03 00:28:55,691 Epoch 485/500
+2024-11-03 00:31:55,121 Train Loss: 0.0369893, Val Loss: 0.0562180
+2024-11-03 00:31:55,121 Epoch 486/500
+2024-11-03 00:34:56,180 Train Loss: 0.0372507, Val Loss: 0.0560992
+2024-11-03 00:34:56,180 Epoch 487/500
+2024-11-03 00:38:46,146 Train Loss: 0.0373390, Val Loss: 0.0555817
+2024-11-03 00:38:46,146 Epoch 488/500
+2024-11-03 00:41:47,966 Train Loss: 0.0371835, Val Loss: 0.0561714
+2024-11-03 00:41:47,967 Epoch 489/500
+2024-11-03 00:44:44,799 Train Loss: 0.0372496, Val Loss: 0.0555808
+2024-11-03 00:44:44,800 Epoch 490/500
+2024-11-03 00:48:23,255 Train Loss: 0.0373195, Val Loss: 0.0555320
+2024-11-03 00:48:23,255 Epoch 491/500
+2024-11-03 00:51:35,216 Train Loss: 0.0376244, Val Loss: 0.0558077
+2024-11-03 00:51:35,217 Epoch 492/500
+2024-11-03 00:54:46,920 Train Loss: 0.0376095, Val Loss: 0.0558452
+2024-11-03 00:54:46,921 Epoch 493/500
+2024-11-03 00:58:30,336 Train Loss: 0.0372585, Val Loss: 0.0557160
+2024-11-03 00:58:30,336 Epoch 494/500
+2024-11-03 01:01:33,647 Train Loss: 0.0376758, Val Loss: 0.0563320
+2024-11-03 01:01:33,648 Epoch 495/500
+2024-11-03 01:05:07,852 Train Loss: 0.0371837, Val Loss: 0.0561041
+2024-11-03 01:05:07,853 Epoch 496/500
+2024-11-03 01:08:25,195 Train Loss: 0.0376199, Val Loss: 0.0559861
+2024-11-03 01:08:25,195 Epoch 497/500
+2024-11-03 01:11:37,813 Train Loss: 0.0373268, Val Loss: 0.0553711
+2024-11-03 01:11:37,814 Epoch 498/500
+2024-11-03 01:15:16,007 Train Loss: 0.0373039, Val Loss: 0.0554337
+2024-11-03 01:15:16,007 Epoch 499/500
+2024-11-03 01:18:15,856 Train Loss: 0.0373263, Val Loss: 0.0555723
+2024-11-03 01:18:15,856 Epoch 500/500
+2024-11-03 01:21:55,890 Train Loss: 0.0373820, Val Loss: 0.0555461
+2024-11-03 01:22:22,918 Testing completed and best model saved.

Exp1_Global_weather_forecasting/model/Triton_model.py

@@ -0,0 +1,516 @@
+import torch
+from torch import nn
+import math
+from timm.layers import DropPath, trunc_normal_
+
+def stride_generator(N, reverse=False):
+    strides = [1, 2] * 10
+    if reverse:
+        return list(reversed(strides[:N]))
+    else:
+        return strides[:N]
+    
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super(MLP, self).__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class ConvMLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super(ConvMLP, self).__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.act = act_layer()
+        self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+class ConvBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=4,
+        mlp_ratio=4.,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.,
+        attn_drop=0.,
+        drop_path=0.,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm
+    ):
+        super(ConvBlock, self).__init__()
+        self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
+        self.norm1 = nn.BatchNorm2d(dim)
+        self.conv1 = nn.Conv2d(dim, dim, 1)
+        self.conv2 = nn.Conv2d(dim, dim, 1)
+        self.attn = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = nn.BatchNorm2d(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ConvMLP(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop
+        )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = (
+                m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            )
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {}
+
+    def forward(self, x):
+        x = x + self.pos_embed(x)
+        x = x + self.drop_path(
+            self.conv2(self.attn(self.conv1(self.norm1(x))))
+        )
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.,
+        attn_drop=0.,
+        drop_path=0.,
+        init_value=1e-6,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm
+    ):
+        super(SelfAttentionBlock, self).__init__()
+        self.pos_embed = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop
+        )
+        self.gamma_1 = nn.Parameter(init_value * torch.ones((dim)), requires_grad=True)
+        self.gamma_2 = nn.Parameter(init_value * torch.ones((dim)), requires_grad=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'gamma_1', 'gamma_2'}
+
+    def forward(self, x):
+        x = x + self.pos_embed(x)
+        B, N, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        x = x.transpose(1, 2).reshape(B, N, H, W)
+        return x
+
+def UniformerSubBlock(
+    embed_dims,
+    mlp_ratio=4.,
+    drop=0.,
+    drop_path=0.,
+    init_value=1e-6,
+    block_type='Conv'
+):
+    assert block_type in ['Conv', 'MHSA']
+    if block_type == 'Conv':
+        # return ConvBlock(dim=embed_dims, mlp_ratio=mlp_ratio, drop=drop, drop_path=drop_path)
+        return SelfAttentionBlock(
+            dim=embed_dims,
+            num_heads=8,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=True,
+            drop=drop,
+            drop_path=drop_path,
+            init_value=init_value
+        )
+    else:
+        return SelfAttentionBlock(
+            dim=embed_dims,
+            num_heads=8,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=True,
+            drop=drop,
+            drop_path=drop_path,
+            init_value=init_value
+        )
+
+class SpatioTemporalEvolutionBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        input_resolution=None,
+        mlp_ratio=8.,
+        drop=0.0,
+        drop_path=0.0,
+        layer_i=0
+    ):
+        super(SpatioTemporalEvolutionBlock, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        block_type = 'MHSA' if in_channels == out_channels and layer_i > 0 else 'Conv'
+        self.block = UniformerSubBlock(
+            in_channels,
+            mlp_ratio=mlp_ratio,
+            drop=drop,
+            drop_path=drop_path,
+            block_type=block_type
+        )
+
+        if in_channels != out_channels:
+            self.reduction = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0
+            )
+
+    def forward(self, x):
+        z = self.block(x)
+        if self.in_channels != self.out_channels:
+            z = self.reduction(z)
+        return z
+
+class SpatioTemporalEvolution(nn.Module):
+    def __init__(
+        self,
+        channel_in,
+        channel_hid,
+        N2,
+        input_resolution=None,
+        mlp_ratio=4.,
+        drop=0.0,
+        drop_path=0.1
+    ):
+        super(SpatioTemporalEvolution, self).__init__()
+        assert N2 >= 2 and mlp_ratio > 1
+        self.N2 = N2
+        dpr = [x.item() for x in torch.linspace(1e-2, drop_path, self.N2)]
+
+        evolution_layers = [SpatioTemporalEvolutionBlock(
+            channel_in,
+            channel_hid,
+            input_resolution,
+            mlp_ratio=mlp_ratio,
+            drop=drop,
+            drop_path=dpr[0],
+            layer_i=0
+        )]
+
+        for i in range(1, N2 - 1):
+            evolution_layers.append(SpatioTemporalEvolutionBlock(
+                channel_hid,
+                channel_hid,
+                input_resolution,
+                mlp_ratio=mlp_ratio,
+                drop=drop,
+                drop_path=dpr[i],
+                layer_i=i
+            ))
+
+        evolution_layers.append(SpatioTemporalEvolutionBlock(
+            channel_hid,
+            channel_in,
+            input_resolution,
+            mlp_ratio=mlp_ratio,
+            drop=drop,
+            drop_path=drop_path,
+            layer_i=N2 - 1
+        ))
+        self.enc = nn.Sequential(*evolution_layers)
+
+    def forward(self, x):
+        B, T, C, H, W = x.shape
+        x = x.reshape(B, T * C, H, W)
+        z = x
+        for i in range(self.N2):
+            z = self.enc[i](z)
+        y = z.reshape(B, T, C, H, W)
+        return y
+
+class BasicConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        transpose=False,
+        act_norm=False
+    ):
+        super(BasicConv2d, self).__init__()
+        self.act_norm = act_norm
+        if not transpose:
+            self.conv = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding
+            )
+        else:
+            self.conv = nn.ConvTranspose2d(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                output_padding=stride // 2
+            )
+        self.norm = nn.GroupNorm(2, out_channels)
+        self.act = nn.LeakyReLU(0.2, inplace=True)
+
+    def forward(self, x):
+        y = self.conv(x)
+        if self.act_norm:
+            y = self.act(self.norm(y))
+        return y
+
+class ConvDynamicsLayer(nn.Module):
+    def __init__(self, C_in, C_out, stride, transpose=False, act_norm=True):
+        super(ConvDynamicsLayer, self).__init__()
+        if stride == 1:
+            transpose = False
+        self.conv = BasicConv2d(
+            C_in,
+            C_out,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            transpose=transpose,
+            act_norm=act_norm
+        )
+
+    def forward(self, x):
+        y = self.conv(x)
+        return y
+
+class MultiGroupConv2d(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups,
+        act_norm=False
+    ):
+        super(MultiGroupConv2d, self).__init__()
+        self.act_norm = act_norm
+        if in_channels % groups != 0:
+            groups = 1
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups
+        )
+        self.norm = nn.GroupNorm(groups, out_channels)
+        self.activate = nn.LeakyReLU(0.2, inplace=True)
+
+    def forward(self, x):
+        y = self.conv(x)
+        if self.act_norm:
+            y = self.activate(self.norm(y))
+        return y
+
+
+class AtmosphericEncoder(nn.Module):
+    def __init__(self, C_in, spatial_hidden_dim, num_spatial_layers):
+        super(AtmosphericEncoder, self).__init__()
+        strides = stride_generator(num_spatial_layers)
+        self.enc = nn.Sequential(
+            ConvDynamicsLayer(C_in, spatial_hidden_dim, stride=strides[0]),
+            *[ConvDynamicsLayer(spatial_hidden_dim, spatial_hidden_dim, stride=s) for s in strides[1:]]
+        )
+
+    def forward(self, x):
+        enc1 = self.enc[0](x)
+        latent = enc1
+        for i in range(1, len(self.enc)):
+            latent = self.enc[i](latent)
+        return latent, enc1
+
+class AtmosphericDecoder(nn.Module):
+    def __init__(self, spatial_hidden_dim, C_out, num_spatial_layers):
+        super(AtmosphericDecoder, self).__init__()
+        strides = stride_generator(num_spatial_layers, reverse=True)
+        self.dec = nn.Sequential(
+            *[ConvDynamicsLayer(spatial_hidden_dim, spatial_hidden_dim, stride=s, transpose=True) for s in strides[:-1]],
+            ConvDynamicsLayer(2 * spatial_hidden_dim, spatial_hidden_dim, stride=strides[-1], transpose=True)
+        )
+        self.readout = nn.Conv2d(spatial_hidden_dim, C_out, 1)
+
+    def forward(self, hid, enc1=None):
+        for i in range(0, len(self.dec) - 1):
+            hid = self.dec[i](hid)
+        Y = self.dec[-1](torch.cat([hid, enc1], dim=1))
+        Y = self.readout(Y)
+        return Y
+
+class Triton(nn.Module):
+    def __init__(
+        self,
+        shape_in,
+        spatial_hidden_dim=64,
+        output_channels=4,
+        temporal_hidden_dim=128,
+        num_spatial_layers=4,
+        num_temporal_layers=8,
+        in_time_seq_length=10,
+        out_time_seq_length=10
+    ):
+        super(Triton, self).__init__()
+        T, C, H, W = shape_in
+        self.H1 = int(H / 2 ** (num_spatial_layers / 2)) + 1 if H % 3 == 0 else int(H / 2 ** (num_spatial_layers / 2))
+        self.W1 = int(W / 2 ** (num_spatial_layers / 2))
+        self.output_dim = output_channels
+        self.input_time_seq_length = in_time_seq_length
+        self.output_time_seq_length = out_time_seq_length
+        
+        self.atmospheric_encoder = AtmosphericEncoder(C, spatial_hidden_dim, num_spatial_layers)
+        self.temporal_evolution = SpatioTemporalEvolution(
+            T * spatial_hidden_dim,
+            temporal_hidden_dim,
+            num_temporal_layers,
+            input_resolution=[self.H1, self.W1],
+            mlp_ratio=4.0,
+            drop_path=0.1
+        )
+        self.atmospheric_decoder = AtmosphericDecoder(spatial_hidden_dim, self.output_dim, num_spatial_layers)
+
+    def forward(self, input_state):
+        """
+        1. Reshape the input state to match the encoder's input requirements.
+        2. Extract features using the Atmospheric Encoder and obtain skip connections.
+        3. Perform spatio-temporal evolution on the encoded features.
+        4. Decode the evolved features to generate the final output.
+        """
+        batch_size, temporal_length, channels, height, width = input_state.shape
+        reshaped_input = input_state.view(batch_size * temporal_length, channels, height, width)
+        
+        encoded_features, skip_connection = self.atmospheric_encoder(reshaped_input)
+        _, encoded_channels, encoded_height, encoded_width = encoded_features.shape
+        encoded_features = encoded_features.view(batch_size, temporal_length, encoded_channels, encoded_height, encoded_width)
+        
+        temporal_bias = encoded_features
+        temporal_hidden = self.temporal_evolution(temporal_bias)
+        reshaped_hidden = temporal_hidden.view(batch_size * temporal_length, encoded_channels, encoded_height, encoded_width)
+
+        decoded_output = self.atmospheric_decoder(reshaped_hidden, skip_connection)
+        final_output = decoded_output.view(batch_size, temporal_length, -1, height, width)
+        
+        return final_output
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+if __name__ == '__main__':
+    inputs = torch.randn(1, 1, 69, 180, 360)
+    model = Triton(
+        shape_in=(1, 69, 180, 360),
+        spatial_hidden_dim=64,
+        output_channels=69,
+        temporal_hidden_dim=128,
+        num_spatial_layers=4,
+        num_temporal_layers=8)
+    output = model(inputs)
+    print(output.shape)

Exp1_Global_weather_forecasting/model_baselines/fuxi_model.py

@@ -0,0 +1,242 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from timm.layers.helpers import to_2tuple
+from timm.models.swin_transformer_v2 import SwinTransformerV2Stage
+
+from typing import Sequence
+
+
+
+def get_pad3d(input_resolution, window_size):
+    """
+    Args:
+        input_resolution (tuple[int]): (Pl, Lat, Lon)
+        window_size (tuple[int]): (Pl, Lat, Lon)
+
+    Returns:
+        padding (tuple[int]): (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)
+    """
+    Pl, Lat, Lon = input_resolution
+    win_pl, win_lat, win_lon = window_size
+
+    padding_left = padding_right = padding_top = padding_bottom = padding_front = padding_back = 0
+    pl_remainder = Pl % win_pl
+    lat_remainder = Lat % win_lat
+    lon_remainder = Lon % win_lon
+
+    if pl_remainder:
+        pl_pad = win_pl - pl_remainder
+        padding_front = pl_pad // 2
+        padding_back = pl_pad - padding_front
+    if lat_remainder:
+        lat_pad = win_lat - lat_remainder
+        padding_top = lat_pad // 2
+        padding_bottom = lat_pad - padding_top
+    if lon_remainder:
+        lon_pad = win_lon - lon_remainder
+        padding_left = lon_pad // 2
+        padding_right = lon_pad - padding_left
+
+    return padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back
+
+def get_pad2d(input_resolution, window_size):
+    """
+    Args:
+        input_resolution (tuple[int]): Lat, Lon
+        window_size (tuple[int]): Lat, Lon
+
+    Returns:
+        padding (tuple[int]): (padding_left, padding_right, padding_top, padding_bottom)
+    """
+    input_resolution = [2] + list(input_resolution)
+    window_size = [2] + list(window_size)
+    padding = get_pad3d(input_resolution, window_size)
+    return padding[: 4]
+
+
+class CubeEmbedding(nn.Module):
+    """
+    Args:
+        img_size: T, Lat, Lon
+        patch_size: T, Lat, Lon
+    """
+    def __init__(self, img_size, patch_size, in_chans, embed_dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1], img_size[2] // patch_size[2]]
+
+        self.img_size = img_size
+        self.patches_resolution = patches_resolution
+        self.embed_dim = embed_dim
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x: torch.Tensor):
+        B, C, T, Lat, Lon = x.shape
+        assert T == self.img_size[0] and Lat == self.img_size[1] and Lon == self.img_size[2], \
+            f"Input image size ({T}*{Lat}*{Lon}) doesn't match model ({self.img_size[0]}*{self.img_size[1]}*{self.img_size[2]})."
+        x = self.proj(x).reshape(B, self.embed_dim, -1).transpose(1, 2)  # B T*Lat*Lon C
+        if self.norm is not None:
+            x = self.norm(x)
+        x = x.transpose(1, 2).reshape(B, self.embed_dim, *self.patches_resolution)
+        return x
+
+
+class DownBlock(nn.Module):
+    def __init__(self, in_chans: int, out_chans: int, num_groups: int, num_residuals: int = 2):
+        super().__init__()
+        self.conv = nn.Conv2d(in_chans, out_chans, kernel_size=(3, 3), stride=2, padding=1)
+
+        blk = []
+        for i in range(num_residuals):
+            blk.append(nn.Conv2d(out_chans, out_chans, kernel_size=3, stride=1, padding=1))
+            blk.append(nn.GroupNorm(num_groups, out_chans))
+            blk.append(nn.SiLU())
+
+        self.b = nn.Sequential(*blk)
+
+    def forward(self, x):
+        _, _, h, w = x.shape
+        x = self.conv(x)
+
+        shortcut = x
+
+        x = self.b(x)
+
+        res = x + shortcut
+        if h % 2 != 0:
+            res = res[:, :, :-1, :]
+        if w % 2 != 0:
+            res = res[:, :, :, :-1]
+        return res
+
+
+class UpBlock(nn.Module):
+    def __init__(self, in_chans, out_chans, num_groups, num_residuals=2):
+        super().__init__()
+        self.conv = nn.ConvTranspose2d(in_chans, out_chans, kernel_size=2, stride=2)
+
+        blk = []
+        for i in range(num_residuals):
+            blk.append(nn.Conv2d(out_chans, out_chans, kernel_size=3, stride=1, padding=1))
+            blk.append(nn.GroupNorm(num_groups, out_chans))
+            blk.append(nn.SiLU())
+
+        self.b = nn.Sequential(*blk)
+
+    def forward(self, x):
+        x = self.conv(x)
+
+        shortcut = x
+
+        x = self.b(x)
+
+        return x + shortcut
+
+
+class UTransformer(nn.Module):
+    """U-Transformer
+    Args:
+        embed_dim (int): Patch embedding dimension.
+        num_groups (int | tuple[int]): number of groups to separate the channels into.
+        input_resolution (tuple[int]): Lat, Lon.
+        num_heads (int): Number of attention heads in different layers.
+        window_size (int | tuple[int]): Window size.
+        depth (int): Number of blocks.
+    """
+    def __init__(self, embed_dim, num_groups, input_resolution, num_heads, window_size, depth):
+        super().__init__()
+        num_groups = to_2tuple(num_groups)
+        window_size = to_2tuple(window_size)
+        padding = get_pad2d(input_resolution, window_size)
+        padding_left, padding_right, padding_top, padding_bottom = padding
+        self.padding = padding
+        self.pad = nn.ZeroPad2d(padding)
+        input_resolution = list(input_resolution)
+        input_resolution[0] = input_resolution[0] + padding_top + padding_bottom
+        input_resolution[1] = input_resolution[1] + padding_left + padding_right
+        self.down = DownBlock(embed_dim, embed_dim, num_groups[0])
+        self.layer = SwinTransformerV2Stage(embed_dim, embed_dim, input_resolution, depth, num_heads, window_size)
+        self.up = UpBlock(embed_dim * 2, embed_dim, num_groups[1])
+
+    def forward(self, x):
+        B, C, Lat, Lon = x.shape
+        padding_left, padding_right, padding_top, padding_bottom = self.padding
+        x = self.down(x)
+
+        shortcut = x
+
+        # pad
+        x = self.pad(x)
+        _, _, pad_lat, pad_lon = x.shape
+
+        x = x.permute(0, 2, 3, 1)  # B Lat Lon C
+        x = self.layer(x)
+        x = x.permute(0, 3, 1, 2)
+
+        # crop
+        x = x[:, :, padding_top: pad_lat - padding_bottom, padding_left: pad_lon - padding_right]
+
+        # concat
+        x = torch.cat([shortcut, x], dim=1)  # B 2*C Lat Lon
+
+        x = self.up(x)
+        return x
+
+
+class Fuxi(nn.Module):
+    """
+    Args:
+        img_size (Sequence[int], optional): T, Lat, Lon.
+        patch_size (Sequence[int], optional): T, Lat, Lon.
+        in_chans (int, optional): number of input channels.
+        out_chans (int, optional): number of output channels.
+        embed_dim (int, optional): number of embed channels.
+        num_groups (Sequence[int] | int, optional): number of groups to separate the channels into.
+        num_heads (int, optional): Number of attention heads.
+        window_size (int | tuple[int], optional): Local window size.
+    """
+    def __init__(self, in_shape=(1, 69, 180, 360), patch_size=(1, 4, 4), in_chans=69, out_chans=69,
+                 embed_dim=1024, num_groups=32, num_heads=8, window_size=7,**kwargs):
+        super().__init__()
+        
+        img_size=in_shape[0],in_shape[2],in_shape[3]
+        input_resolution = int(img_size[1] / patch_size[1] / 2), int(img_size[2] / patch_size[2] / 2)
+        self.cube_embedding = CubeEmbedding(img_size, patch_size, in_chans, embed_dim)
+        self.u_transformer = UTransformer(embed_dim, num_groups, input_resolution, num_heads, window_size, depth=48)
+        self.fc = nn.Linear(embed_dim, out_chans * patch_size[1] * patch_size[2])
+
+        self.patch_size = patch_size
+        self.input_resolution = input_resolution
+        self.out_chans = out_chans
+        self.img_size = img_size
+
+    def forward(self, x: torch.Tensor):
+        x = x.permute(0, 2, 1, 3, 4)
+        B, _, _, _, _ = x.shape
+        _, patch_lat, patch_lon = self.patch_size
+        Lat, Lon = self.input_resolution
+        Lat, Lon = Lat * 2, Lon * 2
+        x = self.cube_embedding(x).squeeze(2)  # B C Lat Lon
+        x = self.u_transformer(x)
+        x = self.fc(x.permute(0, 2, 3, 1))  # B Lat Lon C
+        x = x.reshape(B, Lat, Lon, patch_lat, patch_lon, self.out_chans).permute(0, 1, 3, 2, 4, 5)
+        # B, lat, patch_lat, lon, patch_lon, C
+
+        x = x.reshape(B, Lat * patch_lat, Lon * patch_lon, self.out_chans)
+        x = x.permute(0, 3, 1, 2)  # B C Lat Lon
+
+        # bilinear
+        x = F.interpolate(x, size=self.img_size[1:], mode="bilinear", align_corners=True).unsqueeze(1)
+    
+        return x
+    
+if __name__ == '__main__':
+    inputs = torch.randn(1, 1, 69, 180, 360)
+    model = Fuxi(in_shape=(1, 69, 180, 360)) #in_shape=(1, 69, 180, 360)
+    output = model(inputs)
+    print(inputs.shape)
+    print(output.shape)

Exp1_Global_weather_forecasting/model_baselines/pangu_model.py

@@ -0,0 +1,1218 @@
+import torch
+from torch import nn
+import numpy as np
+import math
+from collections.abc import Sequence
+import warnings
+
+
+
+##### weight init ######
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    u1 = norm_cdf((a - mean) / std)
+    u2 = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [u1, u2], then translate to
+    # [2u1-1, 2u2-1].
+    tensor.uniform_(2 * u1 - 1, 2 * u2 - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.0))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+    return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Cut & paste from timm master
+    Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    NOTE: this impl is similar to the PyTorch trunc_normal_, the bounds [a, b] are
+    applied while sampling the normal with mean/std applied, therefore a, b args
+    should be adjusted to match the range of mean, std args.
+    """
+    with torch.no_grad():
+        return _trunc_normal_(tensor, mean, std, a, b)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+
+def drop_path(
+    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+):
+    """Cut & paste from timm master
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Cut & paste from timm master
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f"drop_prob={round(self.drop_prob,3):0.3f}"
+
+
+
+class PatchEmbed2D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    2D Image to Patch Embedding.
+
+    Args:
+        img_size (tuple[int]): Image size.
+        patch_size (tuple[int]): Patch token size.
+        in_chans (int): Number of input image channels.
+        embed_dim(int): Number of projection output channels.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size, patch_size, in_chans, embed_dim, norm_layer=None):
+        super().__init__()
+        self.img_size = img_size
+        height, width = img_size
+        h_patch_size, w_path_size = patch_size
+        padding_left = padding_right = padding_top = padding_bottom = 0
+
+        h_remainder = height % h_patch_size
+        w_remainder = width % w_path_size
+
+        if h_remainder:
+            h_pad = h_patch_size - h_remainder
+            padding_top = h_pad // 2
+            padding_bottom = int(h_pad - padding_top)
+
+        if w_remainder:
+            w_pad = w_path_size - w_remainder
+            padding_left = w_pad // 2
+            padding_right = int(w_pad - padding_left)
+
+        self.pad = nn.ZeroPad2d(
+            (padding_left, padding_right, padding_top, padding_bottom)
+        )
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x: torch.Tensor):
+        B, C, H, W = x.shape
+        x = self.pad(x)
+        x = self.proj(x)
+        if self.norm is not None:
+            x = self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        return x
+
+
+class PatchEmbed3D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    3D Image to Patch Embedding.
+
+    Args:
+        img_size (tuple[int]): Image size.
+        patch_size (tuple[int]): Patch token size.
+        in_chans (int): Number of input image channels.
+        embed_dim(int): Number of projection output channels.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size, patch_size, in_chans, embed_dim, norm_layer=None):
+        super().__init__()
+        self.img_size = img_size
+        level, height, width = img_size
+        l_patch_size, h_patch_size, w_patch_size = patch_size
+        padding_left = (
+            padding_right
+        ) = padding_top = padding_bottom = padding_front = padding_back = 0
+
+        l_remainder = level % l_patch_size
+        h_remainder = height % l_patch_size
+        w_remainder = width % w_patch_size
+
+        if l_remainder:
+            l_pad = l_patch_size - l_remainder
+            padding_front = l_pad // 2
+            padding_back = l_pad - padding_front
+        if h_remainder:
+            h_pad = h_patch_size - h_remainder
+            padding_top = h_pad // 2
+            padding_bottom = h_pad - padding_top
+        if w_remainder:
+            w_pad = w_patch_size - w_remainder
+            padding_left = w_pad // 2
+            padding_right = w_pad - padding_left
+
+        self.pad = nn.ConstantPad3d(
+    (
+        padding_left,
+        padding_right,
+        padding_top,
+        padding_bottom,
+        padding_front,
+        padding_back,
+    ),
+    value=0  
+)
+        self.proj = nn.Conv3d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x: torch.Tensor):
+        B, C, L, H, W = x.shape
+        x = self.pad(x)
+        x = self.proj(x)
+        if self.norm:
+            x = self.norm(x.permute(0, 2, 3, 4, 1)).permute(0, 4, 1, 2, 3)
+        return x
+
+
+class PatchRecovery2D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    Patch Embedding Recovery to 2D Image.
+
+    Args:
+        img_size (tuple[int]): Lat, Lon
+        patch_size (tuple[int]): Lat, Lon
+        in_chans (int): Number of input channels.
+        out_chans (int): Number of output channels.
+    """
+
+    def __init__(self, img_size, patch_size, in_chans, out_chans):
+        super().__init__()
+        self.img_size = img_size
+        self.conv = nn.ConvTranspose2d(in_chans, out_chans, patch_size, patch_size)
+
+    def forward(self, x):
+        output = self.conv(x)
+        _, _, H, W = output.shape
+        h_pad = H - self.img_size[0]
+        w_pad = W - self.img_size[1]
+
+        padding_top = h_pad // 2
+        padding_bottom = int(h_pad - padding_top)
+
+        padding_left = w_pad // 2
+        padding_right = int(w_pad - padding_left)
+
+        return output[
+            :, :, padding_top : H - padding_bottom, padding_left : W - padding_right
+        ]
+
+
+class PatchRecovery3D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    Patch Embedding Recovery to 3D Image.
+
+    Args:
+        img_size (tuple[int]): Pl, Lat, Lon
+        patch_size (tuple[int]): Pl, Lat, Lon
+        in_chans (int): Number of input channels.
+        out_chans (int): Number of output channels.
+    """
+
+    def __init__(self, img_size, patch_size, in_chans, out_chans):
+        super().__init__()
+        self.img_size = img_size
+        self.conv = nn.ConvTranspose3d(in_chans, out_chans, patch_size, patch_size)
+
+    def forward(self, x: torch.Tensor):
+        output = self.conv(x)
+        _, _, Pl, Lat, Lon = output.shape
+
+        pl_pad = Pl - self.img_size[0]
+        lat_pad = Lat - self.img_size[1]
+        lon_pad = Lon - self.img_size[2]
+
+        padding_front = pl_pad // 2
+        padding_back = pl_pad - padding_front
+
+        padding_top = lat_pad // 2
+        padding_bottom = lat_pad - padding_top
+
+        padding_left = lon_pad // 2
+        padding_right = lon_pad - padding_left
+
+        return output[
+            :,
+            :,
+            padding_front : Pl - padding_back,
+            padding_top : Lat - padding_bottom,
+            padding_left : Lon - padding_right,
+        ]
+
+class UpSample3D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    3D Up-sampling operation.
+    Implementation from: https://github.com/198808xc/Pangu-Weather/blob/main/pseudocode.py
+
+    Args:
+        in_dim (int): Number of input channels.
+        out_dim (int): Number of output channels.
+        input_resolution (tuple[int]): [pressure levels, latitude, longitude]
+        output_resolution (tuple[int]): [pressure levels, latitude, longitude]
+    """
+
+    def __init__(self, in_dim, out_dim, input_resolution, output_resolution):
+        super().__init__()
+        self.linear1 = nn.Linear(in_dim, out_dim * 4, bias=False)
+        self.linear2 = nn.Linear(out_dim, out_dim, bias=False)
+        self.norm = nn.LayerNorm(out_dim)
+        self.input_resolution = input_resolution
+        self.output_resolution = output_resolution
+
+    def forward(self, x: torch.Tensor):
+        """
+        Args:
+            x (torch.Tensor): (B, N, C)
+        """
+        B, N, C = x.shape
+        in_pl, in_lat, in_lon = self.input_resolution
+        out_pl, out_lat, out_lon = self.output_resolution
+
+        x = self.linear1(x)
+        x = x.reshape(B, in_pl, in_lat, in_lon, 2, 2, C // 2).permute(
+            0, 1, 2, 4, 3, 5, 6
+        )
+        x = x.reshape(B, in_pl, in_lat * 2, in_lon * 2, -1)
+
+        pad_h = in_lat * 2 - out_lat
+        pad_w = in_lon * 2 - out_lon
+
+        pad_top = pad_h // 2
+        pad_bottom = pad_h - pad_top
+
+        pad_left = pad_w // 2
+        pad_right = pad_w - pad_left
+
+        x = x[
+            :,
+            :out_pl,
+            pad_top : 2 * in_lat - pad_bottom,
+            pad_left : 2 * in_lon - pad_right,
+            :,
+        ]
+        x = x.reshape(x.shape[0], x.shape[1] * x.shape[2] * x.shape[3], x.shape[4])
+        x = self.norm(x)
+        x = self.linear2(x)
+        return x
+
+class DownSample3D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    3D Down-sampling operation
+    Implementation from: https://github.com/198808xc/Pangu-Weather/blob/main/pseudocode.py
+
+    Args:
+        in_dim (int): Number of input channels.
+        input_resolution (tuple[int]): [pressure levels, latitude, longitude]
+        output_resolution (tuple[int]): [pressure levels, latitude, longitude]
+    """
+
+    def __init__(self, in_dim, input_resolution, output_resolution):
+        super().__init__()
+        self.linear = nn.Linear(in_dim * 4, in_dim * 2, bias=False)
+        self.norm = nn.LayerNorm(4 * in_dim)
+        self.input_resolution = input_resolution
+        self.output_resolution = output_resolution
+
+        in_pl, in_lat, in_lon = self.input_resolution
+        out_pl, out_lat, out_lon = self.output_resolution
+
+        h_pad = out_lat * 2 - in_lat
+        w_pad = out_lon * 2 - in_lon
+
+        pad_top = h_pad // 2
+        pad_bottom = h_pad - pad_top
+
+        pad_left = w_pad // 2
+        pad_right = w_pad - pad_left
+
+        pad_front = pad_back = 0
+
+        self.pad = nn.ConstantPad3d(
+            (pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back), value=0
+        )
+
+    def forward(self, x):
+        B, N, C = x.shape
+        in_pl, in_lat, in_lon = self.input_resolution
+        out_pl, out_lat, out_lon = self.output_resolution
+        x = x.reshape(B, in_pl, in_lat, in_lon, C)
+
+        # Padding the input to facilitate downsampling
+        x = self.pad(x.permute(0, -1, 1, 2, 3)).permute(0, 2, 3, 4, 1)
+        x = x.reshape(B, in_pl, out_lat, 2, out_lon, 2, C).permute(0, 1, 2, 4, 3, 5, 6)
+        x = x.reshape(B, out_pl * out_lat * out_lon, 4 * C)
+
+        x = self.norm(x)
+        x = self.linear(x)
+        return x
+
+
+
+def get_earth_position_index(window_size, ndim=3):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    This function construct the position index to reuse symmetrical parameters of the position bias.
+    implementation from: https://github.com/198808xc/Pangu-Weather/blob/main/pseudocode.py
+
+    Args:
+        window_size (tuple[int]): [pressure levels, latitude, longitude] or [latitude, longitude]
+        ndim (int): dimension of tensor, 3 or 2
+
+    Returns:
+        position_index (torch.Tensor): [win_pl * win_lat * win_lon, win_pl * win_lat * win_lon] or [win_lat * win_lon, win_lat * win_lon]
+    """
+    if ndim == 3:
+        win_pl, win_lat, win_lon = window_size
+    elif ndim == 2:
+        win_lat, win_lon = window_size
+
+    if ndim == 3:
+        # Index in the pressure level of query matrix
+        coords_zi = torch.arange(win_pl)
+        # Index in the pressure level of key matrix
+        coords_zj = -torch.arange(win_pl) * win_pl
+
+    # Index in the latitude of query matrix
+    coords_hi = torch.arange(win_lat)
+    # Index in the latitude of key matrix
+    coords_hj = -torch.arange(win_lat) * win_lat
+
+    # Index in the longitude of the key-value pair
+    coords_w = torch.arange(win_lon)
+
+    # Change the order of the index to calculate the index in total
+    if ndim == 3:
+        coords_1 = torch.stack(torch.meshgrid([coords_zi, coords_hi, coords_w]))
+        coords_2 = torch.stack(torch.meshgrid([coords_zj, coords_hj, coords_w]))
+    elif ndim == 2:
+        coords_1 = torch.stack(torch.meshgrid([coords_hi, coords_w]))
+        coords_2 = torch.stack(torch.meshgrid([coords_hj, coords_w]))
+    coords_flatten_1 = torch.flatten(coords_1, 1)
+    coords_flatten_2 = torch.flatten(coords_2, 1)
+    coords = coords_flatten_1[:, :, None] - coords_flatten_2[:, None, :]
+    coords = coords.permute(1, 2, 0).contiguous()
+
+    # Shift the index for each dimension to start from 0
+    if ndim == 3:
+        coords[:, :, 2] += win_lon - 1
+        coords[:, :, 1] *= 2 * win_lon - 1
+        coords[:, :, 0] *= (2 * win_lon - 1) * win_lat * win_lat
+    elif ndim == 2:
+        coords[:, :, 1] += win_lon - 1
+        coords[:, :, 0] *= 2 * win_lon - 1
+
+    # Sum up the indexes in two/three dimensions
+    position_index = coords.sum(-1)
+
+    return position_index
+
+
+def get_pad3d(input_resolution, window_size):
+    """
+    Args:
+        input_resolution (tuple[int]): (Pl, Lat, Lon)
+        window_size (tuple[int]): (Pl, Lat, Lon)
+
+    Returns:
+        padding (tuple[int]): (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)
+    """
+    Pl, Lat, Lon = input_resolution
+    win_pl, win_lat, win_lon = window_size
+
+    padding_left = (
+        padding_right
+    ) = padding_top = padding_bottom = padding_front = padding_back = 0
+    pl_remainder = Pl % win_pl
+    lat_remainder = Lat % win_lat
+    lon_remainder = Lon % win_lon
+
+    if pl_remainder:
+        pl_pad = win_pl - pl_remainder
+        padding_front = pl_pad // 2
+        padding_back = pl_pad - padding_front
+    if lat_remainder:
+        lat_pad = win_lat - lat_remainder
+        padding_top = lat_pad // 2
+        padding_bottom = lat_pad - padding_top
+    if lon_remainder:
+        lon_pad = win_lon - lon_remainder
+        padding_left = lon_pad // 2
+        padding_right = lon_pad - padding_left
+
+    return (
+        padding_left,
+        padding_right,
+        padding_top,
+        padding_bottom,
+        padding_front,
+        padding_back,
+    )
+
+
+
+def crop3d(x: torch.Tensor, resolution):
+    """
+    Args:
+        x (torch.Tensor): B, C, Pl, Lat, Lon
+        resolution (tuple[int]): Pl, Lat, Lon
+    """
+    _, _, Pl, Lat, Lon = x.shape
+    pl_pad = Pl - resolution[0]
+    lat_pad = Lat - resolution[1]
+    lon_pad = Lon - resolution[2]
+
+    padding_front = pl_pad // 2
+    padding_back = pl_pad - padding_front
+
+    padding_top = lat_pad // 2
+    padding_bottom = lat_pad - padding_top
+
+    padding_left = lon_pad // 2
+    padding_right = lon_pad - padding_left
+    return x[
+        :,
+        :,
+        padding_front : Pl - padding_back,
+        padding_top : Lat - padding_bottom,
+        padding_left : Lon - padding_right,
+    ]
+
+
+class EarthAttention3D(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    3D window attention with earth position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): [pressure levels, latitude, longitude]
+        window_size (tuple[int]): [pressure levels, latitude, longitude]
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        window_size,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wpl, Wlat, Wlon
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.type_of_windows = (input_resolution[0] // window_size[0]) * (
+            input_resolution[1] // window_size[1]
+        )
+
+        self.earth_position_bias_table = nn.Parameter(
+            torch.zeros(
+                (window_size[0] ** 2)
+                * (window_size[1] ** 2)
+                * (window_size[2] * 2 - 1),
+                self.type_of_windows,
+                num_heads,
+            )
+        )  # Wpl**2 * Wlat**2 * Wlon*2-1, Npl//Wpl * Nlat//Wlat, nH
+
+        earth_position_index = get_earth_position_index(
+            window_size
+        )  # Wpl*Wlat*Wlon, Wpl*Wlat*Wlon
+        self.register_buffer("earth_position_index", earth_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.earth_position_bias_table = trunc_normal_(
+            self.earth_position_bias_table, std=0.02
+        )
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x: torch.Tensor, mask=None):
+        """
+        Args:
+            x: input features with shape of (B * num_lon, num_pl*num_lat, N, C)
+            mask: (0/-inf) mask with shape of (num_lon, num_pl*num_lat, Wpl*Wlat*Wlon, Wpl*Wlat*Wlon)
+        """
+        B_, nW_, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B_, nW_, N, 3, self.num_heads, C // self.num_heads)
+            .permute(3, 0, 4, 1, 2, 5)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+
+        earth_position_bias = self.earth_position_bias_table[
+            self.earth_position_index.view(-1)
+        ].view(
+            self.window_size[0] * self.window_size[1] * self.window_size[2],
+            self.window_size[0] * self.window_size[1] * self.window_size[2],
+            self.type_of_windows,
+            -1,
+        )  # Wpl*Wlat*Wlon, Wpl*Wlat*Wlon, num_pl*num_lat, nH
+        earth_position_bias = earth_position_bias.permute(
+            3, 2, 0, 1
+        ).contiguous()  # nH, num_pl*num_lat, Wpl*Wlat*Wlon, Wpl*Wlat*Wlon
+        attn = attn + earth_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nLon = mask.shape[0]
+            attn = attn.view(
+                B_ // nLon, nLon, self.num_heads, nW_, N, N
+            ) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, nW_, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).permute(0, 2, 3, 1, 4).reshape(B_, nW_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+
+class Transformer3DBlock(nn.Module):
+    """
+    Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    3D Transformer Block
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (tuple[int]): Window size [pressure levels, latitude, longitude].
+        shift_size (tuple[int]): Shift size for SW-MSA [pressure levels, latitude, longitude].
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        num_heads,
+        window_size=None,
+        shift_size=None,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        window_size = (2, 6, 12) if window_size is None else window_size
+        shift_size = (1, 3, 6) if shift_size is None else shift_size
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+
+        self.norm1 = norm_layer(dim)
+        padding = get_pad3d(input_resolution, window_size)
+        self.pad = nn.ConstantPad3d(padding, value=0)
+
+        pad_resolution = list(input_resolution)
+        pad_resolution[0] += padding[-1] + padding[-2]
+        pad_resolution[1] += padding[2] + padding[3]
+        pad_resolution[2] += padding[0] + padding[1]
+
+        self.attn = EarthAttention3D(
+            dim=dim,
+            input_resolution=pad_resolution,
+            window_size=window_size,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+        shift_pl, shift_lat, shift_lon = self.shift_size
+        self.roll = shift_pl and shift_lon and shift_lat
+
+        if self.roll:
+            attn_mask = get_shift_window_mask(pad_resolution, window_size, shift_size)
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x: torch.Tensor):
+        Pl, Lat, Lon = self.input_resolution
+        B, L, C = x.shape
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, Pl, Lat, Lon, C)
+
+        # start pad
+        x = self.pad(x.permute(0, 4, 1, 2, 3)).permute(0, 2, 3, 4, 1)
+
+        _, Pl_pad, Lat_pad, Lon_pad, _ = x.shape
+
+        shift_pl, shift_lat, shift_lon = self.shift_size
+        if self.roll:
+            shifted_x = torch.roll(
+                x, shifts=(-shift_pl, -shift_lat, -shift_lat), dims=(1, 2, 3)
+            )
+            x_windows = window_partition(shifted_x, self.window_size)
+            # B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C
+        else:
+            shifted_x = x
+            x_windows = window_partition(shifted_x, self.window_size)
+            # B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C
+
+        win_pl, win_lat, win_lon = self.window_size
+        x_windows = x_windows.view(
+            x_windows.shape[0], x_windows.shape[1], win_pl * win_lat * win_lon, C
+        )
+        # B*num_lon, num_pl*num_lat, win_pl*win_lat*win_lon, C
+
+        attn_windows = self.attn(
+            x_windows, mask=self.attn_mask
+        )  # B*num_lon, num_pl*num_lat, win_pl*win_lat*win_lon, C
+
+        attn_windows = attn_windows.view(
+            attn_windows.shape[0], attn_windows.shape[1], win_pl, win_lat, win_lon, C
+        )
+
+        if self.roll:
+            shifted_x = window_reverse(
+                attn_windows, self.window_size, Pl=Pl_pad, Lat=Lat_pad, Lon=Lon_pad
+            )
+            # B * Pl * Lat * Lon * C
+            x = torch.roll(
+                shifted_x, shifts=(shift_pl, shift_lat, shift_lon), dims=(1, 2, 3)
+            )
+        else:
+            shifted_x = window_reverse(
+                attn_windows, self.window_size, Pl=Pl_pad, Lat=Lat_pad, Lon=Lon_pad
+            )
+            x = shifted_x
+
+        # crop, end pad
+        x = crop3d(x.permute(0, 4, 1, 2, 3), self.input_resolution).permute(
+            0, 2, 3, 4, 1
+        )
+
+        x = x.reshape(B, Pl * Lat * Lon, C)
+        x = shortcut + self.drop_path(x)
+
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+
+
+
+##### shift window mask ############
+
+def window_partition(x: torch.Tensor, window_size, ndim=3):
+    """
+    Args:
+        x: (B, Pl, Lat, Lon, C) or (B, Lat, Lon, C)
+        window_size (tuple[int]): [win_pl, win_lat, win_lon] or [win_lat, win_lon]
+        ndim (int): dimension of window (3 or 2)
+
+    Returns:
+        windows: (B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C) or (B*num_lon, num_lat, win_lat, win_lon, C)
+    """
+    if ndim == 3:
+        B, Pl, Lat, Lon, C = x.shape
+        win_pl, win_lat, win_lon = window_size
+        x = x.view(
+            B, Pl // win_pl, win_pl, Lat // win_lat, win_lat, Lon // win_lon, win_lon, C
+        )
+        windows = (
+            x.permute(0, 5, 1, 3, 2, 4, 6, 7)
+            .contiguous()
+            .view(-1, (Pl // win_pl) * (Lat // win_lat), win_pl, win_lat, win_lon, C)
+        )
+        return windows
+    elif ndim == 2:
+        B, Lat, Lon, C = x.shape
+        win_lat, win_lon = window_size
+        x = x.view(B, Lat // win_lat, win_lat, Lon // win_lon, win_lon, C)
+        windows = (
+            x.permute(0, 3, 1, 2, 4, 5)
+            .contiguous()
+            .view(-1, (Lat // win_lat), win_lat, win_lon, C)
+        )
+        return windows
+
+
+def window_reverse(windows, window_size, Pl=1, Lat=1, Lon=1, ndim=3):
+    """
+    Args:
+        windows: (B*num_lon, num_pl*num_lat, win_pl, win_lat, win_lon, C) or (B*num_lon, num_lat, win_lat, win_lon, C)
+        window_size (tuple[int]): [win_pl, win_lat, win_lon] or [win_lat, win_lon]
+        Pl (int): pressure levels
+        Lat (int): latitude
+        Lon (int): longitude
+        ndim (int): dimension of window (3 or 2)
+
+    Returns:
+        x: (B, Pl, Lat, Lon, C) or (B, Lat, Lon, C)
+    """
+    if ndim == 3:
+        win_pl, win_lat, win_lon = window_size
+        B = int(windows.shape[0] / (Lon / win_lon))
+        x = windows.view(
+            B,
+            Lon // win_lon,
+            Pl // win_pl,
+            Lat // win_lat,
+            win_pl,
+            win_lat,
+            win_lon,
+            -1,
+        )
+        x = x.permute(0, 2, 4, 3, 5, 1, 6, 7).contiguous().view(B, Pl, Lat, Lon, -1)
+        return x
+    elif ndim == 2:
+        win_lat, win_lon = window_size
+        B = int(windows.shape[0] / (Lon / win_lon))
+        x = windows.view(B, Lon // win_lon, Lat // win_lat, win_lat, win_lon, -1)
+        x = x.permute(0, 2, 3, 1, 4, 5).contiguous().view(B, Lat, Lon, -1)
+        return x
+
+
+def get_shift_window_mask(input_resolution, window_size, shift_size, ndim=3):
+    """
+    Along the longitude dimension, the leftmost and rightmost indices are actually close to each other.
+    If half windows apper at both leftmost and rightmost positions, they are dircetly merged into one window.
+    Args:
+        input_resolution (tuple[int]): [pressure levels, latitude, longitude] or [latitude, longitude]
+        window_size (tuple[int]): Window size [pressure levels, latitude, longitude] or [latitude, longitude]
+        shift_size (tuple[int]): Shift size for SW-MSA [pressure levels, latitude, longitude] or [latitude, longitude]
+        ndim (int): dimension of window (3 or 2)
+
+    Returns:
+        attn_mask: (n_lon, n_pl*n_lat, win_pl*win_lat*win_lon, win_pl*win_lat*win_lon) or (n_lon, n_lat, win_lat*win_lon, win_lat*win_lon)
+    """
+    if ndim == 3:
+        Pl, Lat, Lon = input_resolution
+        win_pl, win_lat, win_lon = window_size
+        shift_pl, shift_lat, shift_lon = shift_size
+
+        img_mask = torch.zeros((1, Pl, Lat, Lon + shift_lon, 1))
+    elif ndim == 2:
+        Lat, Lon = input_resolution
+        win_lat, win_lon = window_size
+        shift_lat, shift_lon = shift_size
+
+        img_mask = torch.zeros((1, Lat, Lon + shift_lon, 1))
+
+    if ndim == 3:
+        pl_slices = (
+            slice(0, -win_pl),
+            slice(-win_pl, -shift_pl),
+            slice(-shift_pl, None),
+        )
+    lat_slices = (
+        slice(0, -win_lat),
+        slice(-win_lat, -shift_lat),
+        slice(-shift_lat, None),
+    )
+    lon_slices = (
+        slice(0, -win_lon),
+        slice(-win_lon, -shift_lon),
+        slice(-shift_lon, None),
+    )
+
+    cnt = 0
+    if ndim == 3:
+        for pl in pl_slices:
+            for lat in lat_slices:
+                for lon in lon_slices:
+                    img_mask[:, pl, lat, lon, :] = cnt
+                    cnt += 1
+        img_mask = img_mask[:, :, :, :Lon, :]
+    elif ndim == 2:
+        for lat in lat_slices:
+            for lon in lon_slices:
+                img_mask[:, lat, lon, :] = cnt
+                cnt += 1
+        img_mask = img_mask[:, :, :Lon, :]
+
+    mask_windows = window_partition(
+        img_mask, window_size, ndim=ndim
+    )  # n_lon, n_pl*n_lat, win_pl, win_lat, win_lon, 1 or n_lon, n_lat, win_lat, win_lon, 1
+    if ndim == 3:
+        win_total = win_pl * win_lat * win_lon
+    elif ndim == 2:
+        win_total = win_lat * win_lon
+    mask_windows = mask_windows.view(
+        mask_windows.shape[0], mask_windows.shape[1], win_total
+    )
+    attn_mask = mask_windows.unsqueeze(2) - mask_windows.unsqueeze(3)
+    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
+        attn_mask == 0, float(0.0)
+    )
+    return attn_mask
+
+
+
+####### FuserLayer ###########
+
+
+class FuserLayer(nn.Module):
+    """Revise from WeatherLearn https://github.com/lizhuoq/WeatherLearn
+    A basic 3D Transformer layer for one stage
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (tuple[int]): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        depth,
+        num_heads,
+        window_size,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+
+        self.blocks = nn.ModuleList(
+            [
+                Transformer3DBlock(
+                    dim=dim,
+                    input_resolution=input_resolution,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    shift_size=(0, 0, 0) if i % 2 == 0 else None,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=drop_path[i]
+                    if isinstance(drop_path, Sequence)
+                    else drop_path,
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+
+    def forward(self, x):
+        for blk in self.blocks:
+            x = blk(x)
+        return x
+
+
+
+class Pangu(nn.Module):
+    """
+    Pangu A PyTorch impl of: `Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast`
+    - https://arxiv.org/abs/2211.02556
+
+    Args:
+        img_size (tuple[int]): Image size [Lat, Lon].
+        patch_size (tuple[int]): Patch token size [Lat, Lon].
+        embed_dim (int): Patch embedding dimension. Default: 192
+        num_heads (tuple[int]): Number of attention heads in different layers.
+        window_size (tuple[int]): Window size.
+    """
+
+    def __init__(
+        self,
+        in_shape=(1, 69, 180, 360),
+        patch_size=(2, 4, 4),
+        embed_dim=384,
+        num_heads=(6, 12, 12, 6),
+        window_size=(2, 6, 12),
+    **kwargs):
+        super().__init__()
+
+        
+        drop_path = np.linspace(0, 0.2, 8).tolist()
+        T, C, H, W = in_shape
+        img_size = (H,W)
+        
+        # In addition, three constant masks(the topography mask, land-sea mask and soil type mask)
+        self.patchembed2d = PatchEmbed2D(
+            img_size=img_size,
+            patch_size=patch_size[1:],
+            in_chans=4,  # add
+            embed_dim=embed_dim,
+        )
+        self.patchembed3d = PatchEmbed3D(
+            img_size=(13, img_size[0], img_size[1]),
+            patch_size=patch_size,
+            in_chans=5,
+            embed_dim=embed_dim,
+        )
+        patched_inp_shape = (
+            8,
+            math.ceil(img_size[0] / patch_size[1]),
+            math.ceil(img_size[1] / patch_size[2]),
+        )
+
+        self.layer1 = FuserLayer(
+            dim=embed_dim,
+            input_resolution=patched_inp_shape,
+            depth=2,
+            num_heads=num_heads[0],
+            window_size=window_size,
+            drop_path=drop_path[:2],
+        )
+
+        patched_inp_shape_downsample = (
+            8,
+            math.ceil(patched_inp_shape[1] / 2),
+            math.ceil(patched_inp_shape[2] / 2),
+        )
+        self.downsample = DownSample3D(
+            in_dim=embed_dim,
+            input_resolution=patched_inp_shape,
+            output_resolution=patched_inp_shape_downsample,
+        )
+        self.layer2 = FuserLayer(
+            dim=embed_dim * 2,
+            input_resolution=patched_inp_shape_downsample,
+            depth=6,
+            num_heads=num_heads[1],
+            window_size=window_size,
+            drop_path=drop_path[2:],
+        )
+        self.layer3 = FuserLayer(
+            dim=embed_dim * 2,
+            input_resolution=patched_inp_shape_downsample,
+            depth=6,
+            num_heads=num_heads[2],
+            window_size=window_size,
+            drop_path=drop_path[2:],
+        )
+        self.upsample = UpSample3D(
+            embed_dim * 2, embed_dim, patched_inp_shape_downsample, patched_inp_shape
+        )
+        self.layer4 = FuserLayer(
+            dim=embed_dim,
+            input_resolution=patched_inp_shape,
+            depth=2,
+            num_heads=num_heads[3],
+            window_size=window_size,
+            drop_path=drop_path[:2],
+        )
+        # The outputs of the 2nd encoder layer and the 7th decoder layer are catenated along the channel dimension.
+        self.patchrecovery2d = PatchRecovery2D(
+            img_size, patch_size[1:], 2 * embed_dim, 4
+        )
+        self.patchrecovery3d = PatchRecovery3D(
+            (13, img_size[0], img_size[1]), patch_size, 2 * embed_dim, 5
+        )
+
+    def prepare_input(self, surface, surface_mask, upper_air):
+        """Prepares the input to the model in the required shape.
+        Args:
+            surface (torch.Tensor): 2D n_lat=721, n_lon=1440, chans=4.
+            surface_mask (torch.Tensor): 2D n_lat=721, n_lon=1440, chans=3.
+            upper_air (torch.Tensor): 3D n_pl=13, n_lat=721, n_lon=1440, chans=5.
+        """
+        upper_air = upper_air.reshape(
+            upper_air.shape[0], -1, upper_air.shape[3], upper_air.shape[4]
+        )
+        surface_mask = surface_mask.unsqueeze(0).repeat(surface.shape[0], 1, 1, 1)
+        return torch.cat([surface, surface_mask, upper_air], dim=1)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor): [batch,T,4+5*13, lat, lon]
+        """
+        x = x.squeeze(1)
+        
+        surface = x[:, :4, :, :]
+        upper_air = x[:, 4:, :, :].reshape(x.shape[0], 5, 13, x.shape[2], x.shape[3])
+        surface = self.patchembed2d(surface)
+        upper_air = self.patchembed3d(upper_air)
+
+        x = torch.cat([surface.unsqueeze(2), upper_air], dim=2)
+        B, C, Pl, Lat, Lon = x.shape
+    
+        x = x.reshape(B, C, -1).transpose(1, 2)
+
+        x = self.layer1(x)
+
+        skip = x
+
+        x = self.downsample(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.upsample(x)
+        x = self.layer4(x)
+
+        output = torch.cat([x, skip], dim=-1)
+        output = output.transpose(1, 2).reshape(B, -1, Pl, Lat, Lon)
+        output_surface = output[:, :, 0, :, :]
+        output_upper_air = output[:, :, 1:, :, :]
+
+        output_surface = self.patchrecovery2d(output_surface)
+        output_upper_air = self.patchrecovery3d(output_upper_air)
+        
+        B,_,__,Lat,Lon=output_upper_air.shape
+        final_output = torch.cat([output_surface,output_upper_air.reshape(B,5*13,Lat,Lon)],dim=1).unsqueeze(1)
+        return final_output
+
+if __name__ == '__main__':
+    inputs = torch.randn(1, 1, 69, 180, 360)
+    model = Pangu(in_shape=(1, 69, 180, 360))
+    output = model(inputs)
+    print(inputs.shape)
+    print(output.shape)

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/Fuxi_210.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e5acd314deaabad580ec55c01b76a74013bb05c4d8a54a6c22cb0310c60d252
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/Pangu_210.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfd14601a410d87b38e9d0bb64f59d7cd349c23b91d9f46be480ca7a3f57b115
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/SFNO_210.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af3d4c417521c4ddb4d0091fb99ca8904c913fde0c18985577d07abafd630777
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/Triton_210_day.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6e10e76d95de03068ea788557848d4771974009c7d8f18de4f952a8bda84bac
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/groundtruth_210.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f8c3bf1f921af120ad63ccb4db0c2419e76b6e5e90f3cfbf9968b8e716b502d
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2018/initial_input.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e15ee08e3843e75a12fb6238b87e08eddb1e4b6b371197fcf2278ff4c59944a
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_initial_input.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e15ee08e3843e75a12fb6238b87e08eddb1e4b6b371197fcf2278ff4c59944a
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5891e14364685fbc8da90689f9eb4b70abf83bb06fdb4ef54751374c7296e094
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_13.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ff9b99d9545e9f226225524ef015282e6a7101423f442233339813ab5404253
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_179.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fee6f5eac7cd10f65bdd6cb4caa9dd03cd9295c8a5f69339f1da2086f5a8b2a
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0f7f3a83d1f5d1529e2a6cba57242540ccf124b8203bec0f31873852a86be7d
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_29.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52081d5f9a049eb280a5866cd51e917f757938bf09277e2d6acf10fabd8b5c13
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eddd04fe8f3e47240f2e837538d997ceca0a6e6880e7e0845a633b61c9c0a85c
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_6.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:afb04a18c9dc483b35be03e18102f3369e7704502a47c2303294b2296a53a3fa
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_prediction_step_9.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5c08349d81b1337dee41d1ca7391128b4a75d3a858fbbf43efc17d652ed8b97
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_0.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9f297703389c2d684f268a07f1dc48fbdba0534ce4c37a706f76cd182fa5759
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_13.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0495dbfc103c9d2b69f75710bf64948523c59b6b4b34e133e0316fbd339b4270
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_179.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67cc499e88c0bc24290f1550aeaec58df936a3de221473ae5854cfea8f099582
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7bc93d6593855fccdcd5b35df196dfa5758e33b27b5a1a95e2bbb3ae5663744
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_29.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1175796a11b9d12b4f04ad3c959b3b62b46f4b92723c3542c6b3d88dfb54e040
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de92a245cd366e8b42db8f3554b7eafc1ff2bbeba62ddce752f04bc0ae833c5b
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_6.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f8f379e99bc8337b1aa070f36aa58ea4716de1393cd59af08d84ffdfd19d103
+size 17884928

Exp1_Global_weather_forecasting/plt_triton/S1_weather_forecasting/2020_triton_results/baselines_triton_20250326_exp1_true_label_step_9.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57561ae16acafa5b03eeef0ede659513bfeb1fd718de9f9db0dd80fe9a6641c1
+size 17884928

Exp1_Global_weather_forecasting/train.py

@@ -0,0 +1,203 @@
+import os
+import random
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import torch.utils.data as data
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import netCDF4 as nc
+import torchvision.transforms as transforms
+from dataloader_api.dataloader import *
+from tqdm import tqdm
+import logging
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torch
+from torch.utils.data import DataLoader, ConcatDataset
+from torch.utils.data.distributed import DistributedSampler
+from model.Triton_model import *
+from model_baselines.pangu_model import *
+
+# Setup logging
+backbone = 'triton_weather_20250326_v1'
+logging.basicConfig(filename=f'/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/logs/{backbone}_training_log.log', 
+                    level=logging.INFO, 
+                    format='%(asctime)s %(message)s')
+
+# Set a specific seed
+seed = 42
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+set_seed(seed)
+
+# =========================================================================== dist train ========================================================================================================================
+dist.init_process_group(backend='nccl')
+local_rank = int(os.environ['LOCAL_RANK'])
+torch.cuda.set_device(local_rank)
+device = torch.device("cuda", local_rank)
+num_gpus = torch.cuda.device_count()
+
+def reduce_mean(tensor, nprocs):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= nprocs
+    return rt
+
+# ============================================================= data load ===================================================
+args = {
+        'data_path': '/jizhicfs/easyluwu/scaling_law/ft_local/low_res',
+        'ocean_lead_time': 1,
+        'atmosphere_lead_time': 1,
+        'shuffle': True,
+        'variables_input': list(range(69)),
+        'variables_output': list(range(69)),
+        'lon_start': 0,
+        'lat_start': 0,
+        'lon_end': 1440,
+        'lat_end': 720,
+        'ds_factor': 1,
+    }
+
+train_dataset = train_Dataset(args)
+test_dataset = test_Dataset(args)
+
+train_sampler = data.distributed.DistributedSampler(train_dataset)
+train_loader = data.DataLoader(train_dataset,
+                                num_workers=0,
+                                batch_size=1, 
+                                sampler=train_sampler)
+
+test_sampler = data.distributed.DistributedSampler(test_dataset)
+test_loader = data.DataLoader(test_dataset,
+                            num_workers=0,
+                            batch_size=1,
+                            sampler=test_sampler)
+
+for inputs, targets in iter(train_loader):
+    print(inputs.shape, targets.shape)
+    break
+
+# ================================================ model load ===========================================================
+model = Triton(
+    shape_in=(1, 69, 180, 360),
+    spatial_hidden_dim=256,
+    output_channels=69,
+    temporal_hidden_dim=512,
+    num_spatial_layers=4,
+    num_temporal_layers=8)
+
+model = model.to(device)
+model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
+
+# ============================== criterion and optimizer ======================================================
+criterion = nn.MSELoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.2)
+
+# ===========================train val and test ======================================
+def train(model, train_loader, criterion, optimizer, device):
+    model.train()
+    train_loss = 0.0
+    for inputs, targets in tqdm(train_loader, desc="Training", disable=local_rank != 0):
+        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+        train_loss += loss.item() * inputs.size(0)
+    return train_loss / len(train_loader.dataset)
+
+def validate(model, val_loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    with torch.no_grad():
+        for inputs, targets in tqdm(val_loader, desc="Validation", disable=local_rank != 0):
+            inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+            val_loss += loss.item() * inputs.size(0)
+    return val_loss / len(val_loader.dataset)
+
+def test(model, test_loader, criterion, device):
+    path = '/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/results'
+    model.eval()
+    test_loss = 0.0
+    
+    all_inputs = []
+    all_targets = []
+    all_outputs = []
+    i = 0
+    with torch.no_grad():
+        for inputs, targets in tqdm(test_loader, desc="Testing", disable=local_rank != 0):
+            i += 1
+            print(f"{i} : {inputs.shape}")
+            inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
+            outputs = model(inputs)
+            
+            # Convert tensors to numpy arrays and append to lists
+            all_inputs.append(inputs.cpu().numpy())
+            all_targets.append(targets.cpu().numpy())
+            all_outputs.append(outputs.cpu().numpy())
+            
+            loss = criterion(outputs, targets)
+            test_loss += loss.item() * inputs.size(0)
+
+    all_inputs = np.concatenate(all_inputs, axis=0)
+    all_targets = np.concatenate(all_targets, axis=0)
+    all_outputs = np.concatenate(all_outputs, axis=0)
+
+    if local_rank == 0:
+        np.save(f'{path}/{backbone}_inputs.npy', all_inputs)
+        np.save(f'{path}/{backbone}_targets.npy', all_targets)
+        np.save(f'{path}/{backbone}_outputs.npy', all_outputs)
+    print(test_loss)
+    print(len(test_loader.dataset))
+    print(i)
+    return test_loss / len(test_loader.dataset)
+
+num_epochs = 1000
+best_val_loss = float('inf')
+best_model_path = f'/jizhicfs/easyluwu/ocean_project/NPJ_baselines/Exp_0_Weather/checkpoints/{backbone}_best_model.pth'
+
+if local_rank == 0 and os.path.exists(best_model_path):
+    try:
+        logging.info('Loading best model from checkpoint.')
+        checkpoint = torch.load(best_model_path, map_location=device)
+        model.load_state_dict(checkpoint)
+    except Exception as e:
+        logging.error(f'Error loading model checkpoint: {e}')
+
+for epoch in range(num_epochs):
+    if local_rank == 0:
+        logging.info(f'Epoch {epoch + 1}/{num_epochs}')
+    train_loss = train(model, train_loader, criterion, optimizer, device)
+    val_loss = validate(model, test_loader, criterion, device)
+
+    if local_rank == 0:
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(model.state_dict(), best_model_path)
+        
+        logging.info(f'Train Loss: {train_loss * num_gpus:.7f}, Val Loss: {val_loss * num_gpus:.7f}') 
+
+if local_rank == 0:
+    try:
+        model.load_state_dict(torch.load(best_model_path))
+        test_loss = test(model, test_loader, criterion, device)
+        logging.info(f"Testing completed and best model saved. | test_loss:{test_loss * num_gpus:.7f}")
+    except Exception as e:
+        logging.error(f'Error loading model checkpoint during testing: {e}')
+
+dist.destroy_process_group()