Upload 205 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/current_animation-checkpoint.gif filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/input_output_animation-checkpoint.gif filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/kuroshio_animation-checkpoint.gif filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/ocean_currents_animation-checkpoint.gif filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/sample_animation-checkpoint.gif filter=lfs diff=lfs merge=lfs -text
+Exp3_Kuroshio_forecasting/plt_triton/nmi_vis.ipynb filter=lfs diff=lfs merge=lfs -text

Exp3_Kuroshio_forecasting/checkpoints/Kuro_ConvLSTM_exp1_20250311_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp1_20250221_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp1_20250222_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp1_20250223_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp1_20250224_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp2_20250224_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Dit_exp2_20250316_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Kno_exp1_20250226_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Kno_exp2_20250225_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Simvp_exp1_20250224_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Simvp_exp_128_20250324_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_K_uv_20250218_exp1_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_K_uv_20250218_exp2_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_128_20250322_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_20250221_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_20250222_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_20250224_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_20250224_best_model_prediction.h5

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp1_64_20250323_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp2_20241107_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp3_20241107_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp3_20241111_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_exp_20241107_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Triton_multi_finetune_20250227_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_U_net_exp1_20250225_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_U_net_exp1_20250226_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_U_net_exp2_20250226_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Kuro_Unet_exp_128_20250324_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Triton_Gulf_uv_20250218_exp1_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/Triton_Kuroshio_uv_20250218_exp1_best_model.pth

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

Exp3_Kuroshio_forecasting/checkpoints/dit_kuro_256_20250227_best_model.pth

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

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/Untitled-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader-checkpoint.ipynb

@@ -0,0 +1,397 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "f3b16ba8-ad82-45c1-8119-b6c61e7311b8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Input shape: torch.Size([32, 10, 2, 128, 128])\n",
+      "Target shape: torch.Size([32, 5, 2, 128, 128])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import h5py\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "\n",
+    "class KuroshioDataset(Dataset):\n",
+    "    def __init__(self, data, input_length, target_length, downsample_factor=1):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            data: Tensor of shape (num_samples, num_timesteps, C, H, W)\n",
+    "            input_length: Number of input time steps (T_in)\n",
+    "            target_length: Number of prediction time steps (T_out)\n",
+    "            downsample_factor: Spatial downsampling factor\n",
+    "        \"\"\"\n",
+    "        super().__init__()\n",
+    "        self.data = data\n",
+    "        self.input_length = input_length\n",
+    "        self.target_length = target_length\n",
+    "        self.downsample_factor = downsample_factor\n",
+    "\n",
+    "        # Validate time dimensions\n",
+    "        self.num_samples, self.num_timesteps, self.C, self.H, self.W = data.shape\n",
+    "        self.max_t_start = self.num_timesteps - self.input_length - self.target_length\n",
+    "        assert self.max_t_start >= 0, \"Not enough timesteps for input and output\"\n",
+    "\n",
+    "        # Generate sample indices (sample_idx, t_start)\n",
+    "        self.sample_indices = []\n",
+    "        for s in range(self.num_samples):\n",
+    "            for t_start in range(self.max_t_start + 1):\n",
+    "                self.sample_indices.append((s, t_start))\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.sample_indices)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        s, t_start = self.sample_indices[idx]\n",
+    "        \n",
+    "        # Extract sequences\n",
+    "        input_end = t_start + self.input_length\n",
+    "        output_end = input_end + self.target_length\n",
+    "        \n",
+    "        input_seq = self.data[s, t_start:input_end]      # (T_in, C, H, W)\n",
+    "        target_seq = self.data[s, input_end:output_end]  # (T_out, C, H, W)\n",
+    "\n",
+    "        # Apply downsampling\n",
+    "        if self.downsample_factor > 1:\n",
+    "            dsf = self.downsample_factor\n",
+    "            input_seq = input_seq[..., ::dsf, ::dsf]\n",
+    "            target_seq = target_seq[..., ::dsf, ::dsf]\n",
+    "\n",
+    "        return input_seq.float(), target_seq.float()\n",
+    "\n",
+    "def load_datasets(file_path, args):\n",
+    "    # Load and preprocess data\n",
+    "    with h5py.File(file_path, 'r') as f:\n",
+    "        u_k = np.transpose(f['u_k'][:], (0, 3, 1, 2))  # (2046, 50, 128, 128)\n",
+    "        v_k = np.transpose(f['v_k'][:], (0, 3, 1, 2))\n",
+    "    \n",
+    "    # Combine u and v channels\n",
+    "    combined = np.stack([u_k, v_k], axis=2)  # (2046, 50, 2, 128, 128)\n",
+    "    data_tensor = torch.tensor(combined, dtype=torch.float32)\n",
+    "\n",
+    "    # Split dataset\n",
+    "    total_samples = 2046\n",
+    "    train_size = int(0.8 * total_samples)\n",
+    "    val_size = int(0.1 * total_samples)\n",
+    "    \n",
+    "    train_data = data_tensor[:train_size]\n",
+    "    val_data = data_tensor[train_size:train_size+val_size]\n",
+    "    test_data = data_tensor[train_size+val_size:]\n",
+    "\n",
+    "    # Create datasets\n",
+    "    train_dataset = KuroshioDataset(train_data, \n",
+    "                                  args['input_length'],\n",
+    "                                  args['target_length'],\n",
+    "                                  args['downsample_factor'])\n",
+    "    \n",
+    "    val_dataset = KuroshioDataset(val_data,\n",
+    "                                args['input_length'],\n",
+    "                                args['target_length'],\n",
+    "                                args['downsample_factor'])\n",
+    "    \n",
+    "    test_dataset = KuroshioDataset(test_data,\n",
+    "                                 args['input_length'],\n",
+    "                                 args['target_length'],\n",
+    "                                 args['downsample_factor'])\n",
+    "\n",
+    "    return train_dataset, val_dataset, test_dataset\n",
+    "\n",
+    "# Example usage\n",
+    "if __name__ == \"__main__\":\n",
+    "    config = {\n",
+    "        'input_length': 10,    # T_in: 输入时间步数\n",
+    "        'target_length': 5,    # T_out: 预测时间步数\n",
+    "        'downsample_factor': 1 # 空间下采样因子\n",
+    "    }\n",
+    "\n",
+    "    # 加载数据集\n",
+    "    train_ds, val_ds, test_ds = load_datasets('./Kuroshio_window_data.h5', config)\n",
+    "\n",
+    "    # 创建DataLoader\n",
+    "    batch_size = 32\n",
+    "    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
+    "    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)\n",
+    "    test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False)\n",
+    "\n",
+    "    # 验证数据形状\n",
+    "    sample_input, sample_target = next(iter(train_loader))\n",
+    "    print(f\"Input shape: {sample_input.shape}\")   # 应为 (B, T_in, 2, H, W)\n",
+    "    print(f\"Target shape: {sample_target.shape}\") # 应为 (B, T_out, 2, H, W)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "9c6b1e5c-7874-49f2-9004-c17470a3ae85",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "可视化已保存为 kuroshio_animation.gif\n"
+     ]
+    }
+   ],
+   "source": [
+    "import h5py\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.animation as animation\n",
+    "from matplotlib import gridspec\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "\n",
+    "\n",
+    "# 修正后的可视化函数\n",
+    "def create_visualization(input_seq, target_seq, sample_idx=0, downsample=4, fps=10):\n",
+    "    # 数据准备\n",
+    "    input_np = input_seq[sample_idx].cpu().numpy()\n",
+    "    target_np = target_seq[sample_idx].cpu().numpy()\n",
+    "    full_seq = np.concatenate([input_np, target_np], axis=0)\n",
+    "    full_seq = np.transpose(full_seq, (0, 2, 3, 1))  # [T, H, W, C]\n",
+    "    \n",
+    "    # 创建网格\n",
+    "    H, W = full_seq.shape[1], full_seq.shape[2]\n",
+    "    x = np.arange(W)\n",
+    "    y = np.arange(H)\n",
+    "    X, Y = np.meshgrid(x, y)\n",
+    "    X_ds, Y_ds = X[::downsample, ::downsample], Y[::downsample, ::downsample]\n",
+    "    \n",
+    "    # 计算速度幅值\n",
+    "    speed = np.sqrt(full_seq[...,0]**2 + full_seq[...,1]**2)\n",
+    "    speed_min, speed_max = speed.min(), speed.max()\n",
+    "    \n",
+    "    # 创建画布\n",
+    "    fig = plt.figure(figsize=(15, 5), facecolor='white')\n",
+    "    gs = gridspec.GridSpec(1, 3, width_ratios=[1, 1, 1])\n",
+    "    ax1 = plt.subplot(gs[0])\n",
+    "    ax2 = plt.subplot(gs[1])\n",
+    "    ax3 = plt.subplot(gs[2])\n",
+    "    \n",
+    "    # 初始化子图\n",
+    "    im1 = ax1.imshow(full_seq[0,...,0], origin='lower', cmap='RdBu_r', vmax=1, vmin=-1)\n",
+    "    ax1.set_title(\"U Component\")\n",
+    "    plt.colorbar(im1, ax=ax1)\n",
+    "    \n",
+    "    im2 = ax2.imshow(full_seq[0,...,1], origin='lower', cmap='RdBu_r', vmax=1, vmin=-1)\n",
+    "    ax2.set_title(\"V Component\")\n",
+    "    plt.colorbar(im2, ax=ax2)\n",
+    "    \n",
+    "    # 初始化矢量场\n",
+    "    U = full_seq[0,...,0][::downsample, ::downsample]\n",
+    "    V = full_seq[0,...,1][::downsample, ::downsample]\n",
+    "    speed_initial = np.sqrt(U**2 + V**2)\n",
+    "    quiver = ax3.quiver(X_ds, Y_ds, U, V, speed_initial, \n",
+    "                       cmap='RdBu_r', \n",
+    "                       scale=50, \n",
+    "                       width=0.003,\n",
+    "                       clim=[speed_min, speed_max])\n",
+    "    plt.colorbar(quiver, ax=ax3, label='Flow Speed')\n",
+    "    ax3.set_title(\"Vector Field\")\n",
+    "    \n",
+    "    # 统一设置\n",
+    "    for ax in [ax1, ax2, ax3]:\n",
+    "        ax.set_xticks([])\n",
+    "        ax.set_yticks([])\n",
+    "        ax.set_xlabel(f\"Timestep: 0/{full_seq.shape[0]-1}\")\n",
+    "    \n",
+    "    # 动画更新函数\n",
+    "    def update(frame):\n",
+    "        # 更新分量图\n",
+    "        im1.set_data(full_seq[frame,...,0])\n",
+    "        im2.set_data(full_seq[frame,...,1])\n",
+    "        \n",
+    "        # 更新矢量场\n",
+    "        U = full_seq[frame,...,0][::downsample, ::downsample]\n",
+    "        V = full_seq[frame,...,1][::downsample, ::downsample]\n",
+    "        speed = np.sqrt(U**2 + V**2)\n",
+    "        \n",
+    "        quiver.set_UVC(U, V)\n",
+    "        quiver.set_array(speed.flatten())\n",
+    "        \n",
+    "        # 更新时间标签\n",
+    "        for ax in [ax1, ax2, ax3]:\n",
+    "            ax.set_xlabel(f\"Timestep: {frame}/{full_seq.shape[0]-1}\")\n",
+    "        \n",
+    "        return [im1, im2, quiver]\n",
+    "    \n",
+    "    # 生成动画\n",
+    "    ani = animation.FuncAnimation(fig, update, frames=full_seq.shape[0], \n",
+    "                                interval=1000//fps, blit=False)\n",
+    "    ani.save('kuroshio_animation.gif', writer='pillow', fps=fps)\n",
+    "    plt.close()\n",
+    "    print(\"可视化已保存为 kuroshio_animation.gif\")\n",
+    "\n",
+    "# 完整使用示例\n",
+    "if __name__ == \"__main__\":\n",
+    "    # 配置参数\n",
+    "    config = {\n",
+    "        'input_length': 25,\n",
+    "        'target_length': 25,\n",
+    "        'downsample_factor': 1\n",
+    "    }\n",
+    "    \n",
+    "    # 加载数据\n",
+    "    train_ds, val_ds, test_ds = load_datasets('./Kuroshio_window_data.h5', config)\n",
+    "    train_loader = DataLoader(train_ds, batch_size=10, shuffle=True)\n",
+    "    \n",
+    "    # 获取样本数据\n",
+    "    sample_input, sample_target = next(iter(train_loader))\n",
+    "    \n",
+    "    # 生成可视化（关键参数调整）\n",
+    "    create_visualization(\n",
+    "        sample_input, \n",
+    "        sample_target,\n",
+    "        sample_idx=2,     # 选择样本索引\n",
+    "        downsample=1,     # 矢量场密度（值越小越密集）\n",
+    "        fps=4             # 动画帧率\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d0454a79-3e01-49dd-b4c5-0aca0bd76bcf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import torch.distributed as dist\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torch.utils.data.distributed import DistributedSampler\n",
+    "import h5py\n",
+    "import numpy as np\n",
+    "from torch.utils.data import Dataset\n",
+    "from torch.utils.data import DataLoader\n",
+    "import torchvision.transforms as transforms\n",
+    "import torch.utils.data as data\n",
+    "import h5py\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "class WeatherDataset(Dataset):\n",
+    "    def __init__(self, data_path, horizon, transform=None):\n",
+    "        with h5py.File(data_path, 'r') as f:\n",
+    "            self.data_uv_g = f['u_k'][:]  \n",
+    "            self.data_uv_g = torch.from_numpy(self.data_uv_g).to(torch.float32)\n",
+    "            self.data_uv_g = self.data_uv_g.permute(0, 3, 1, 2).unsqueeze_(2) \n",
+    "            \n",
+    "            self.data_uv_k = f['v_k'][:]  \n",
+    "            self.data_uv_k = torch.from_numpy(self.data_uv_k).to(torch.float32)\n",
+    "            self.data_uv_k = self.data_uv_k.permute(0, 3, 1, 2).unsqueeze_(2) \n",
+    "            self.data_uv_gk = torch.cat([self.data_uv_g, self.data_uv_k], dim=2)\n",
+    "            self.transform = transform\n",
+    "            self.horizon = horizon\n",
+    "            self.mean = 0\n",
+    "            self.std = 1\n",
+    "    \n",
+    "    def __len__(self):\n",
+    "        return len(self.data_uv_gk)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        input_frames = self.data_uv_gk[idx][:self.horizon]\n",
+    "        output_frames = self.data_uv_gk[idx][self.horizon:2*self.horizon]\n",
+    "        input_frames = (input_frames - self.mean) / self.std\n",
+    "        output_frames = (output_frames - self.mean) / self.std\n",
+    "        return input_frames, output_frames\n",
+    "\n",
+    "def load_data(data_path, batch_size, val_batch_size, horizon, num_workers):\n",
+    "    dataset = WeatherDataset(data_path=data_path+'/kg_all_20_mask_latmean.h5', horizon=horizon, transform=None)\n",
+    "    \n",
+    "    total_samples = len(dataset)\n",
+    "    train_size = int(0.8 * total_samples)\n",
+    "    val_size = int(0.1 * total_samples)\n",
+    "    \n",
+    "    train_dataset = dataset[:train_size]\n",
+    "    val_dataset = dataset[train_size:train_size+val_size]\n",
+    "    test_dataset = dataset[train_size+val_size:]\n",
+    "    \n",
+    "    train_sampler = DistributedSampler(train_dataset)\n",
+    "    val_sampler = DistributedSampler(val_dataset)\n",
+    "    test_sampler = DistributedSampler(test_dataset)\n",
+    "\n",
+    "    dataloader_train = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, pin_memory=False,\n",
+    "                                  num_workers=num_workers, drop_last=True)\n",
+    "    dataloader_validation = DataLoader(val_dataset, batch_size=val_batch_size, sampler=val_sampler, pin_memory=False,\n",
+    "                                       num_workers=num_workers, drop_last=True)\n",
+    "    dataloader_test = DataLoader(test_dataset, batch_size=val_batch_size, sampler=test_sampler, pin_memory=False,\n",
+    "                                 num_workers=num_workers, drop_last=True)\n",
+    "    mean, std = 0, 1\n",
+    "\n",
+    "    return dataloader_train, dataloader_validation, dataloader_test, mean, std"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "6b051bb4-492a-4b4a-828a-6099ce9437b4",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'data_tensor' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[3], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[0;32m----> 2\u001b[0m     train_loader, val_loader, test_loader, mean, std \u001b[38;5;241m=\u001b[39m \u001b[43mload_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata_path\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m/jizhicfs/easyluwu/ocean_project/kuro/ft_local\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m                                                                                    \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m                                                                                    \u001b[49m\u001b[43mval_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m                                                                                    \u001b[49m\u001b[43mhorizon\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m                                                                                    \u001b[49m\u001b[43mnum_workers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m input_frames, output_frames \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28miter\u001b[39m(train_loader):\n\u001b[1;32m      8\u001b[0m         \u001b[38;5;28mprint\u001b[39m(input_frames\u001b[38;5;241m.\u001b[39mshape, output_frames\u001b[38;5;241m.\u001b[39mshape)  \u001b[38;5;66;03m# [B, T, C, H, W]\u001b[39;00m\n",
+      "Cell \u001b[0;32mIn[2], line 50\u001b[0m, in \u001b[0;36mload_data\u001b[0;34m(data_path, batch_size, val_batch_size, horizon, num_workers)\u001b[0m\n\u001b[1;32m     47\u001b[0m train_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mint\u001b[39m(\u001b[38;5;241m0.8\u001b[39m \u001b[38;5;241m*\u001b[39m total_samples)\n\u001b[1;32m     48\u001b[0m val_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mint\u001b[39m(\u001b[38;5;241m0.1\u001b[39m \u001b[38;5;241m*\u001b[39m total_samples)\n\u001b[0;32m---> 50\u001b[0m train_dataset \u001b[38;5;241m=\u001b[39m \u001b[43mdata_tensor\u001b[49m[:train_size]\n\u001b[1;32m     51\u001b[0m val_dataset \u001b[38;5;241m=\u001b[39m data_tensor[train_size:train_size\u001b[38;5;241m+\u001b[39mval_size]\n\u001b[1;32m     52\u001b[0m test_dataset \u001b[38;5;241m=\u001b[39m data_tensor[train_size\u001b[38;5;241m+\u001b[39mval_size:]\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'data_tensor' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "if __name__ == '__main__':\n",
+    "    train_loader, val_loader, test_loader, mean, std = load_data(data_path='/jizhicfs/easyluwu/ocean_project/kuro/ft_local',\n",
+    "                                                                                    batch_size=8, \n",
+    "                                                                                    val_batch_size=8, \n",
+    "                                                                                    horizon=10,\n",
+    "                                                                                    num_workers=8)\n",
+    "    for input_frames, output_frames in iter(train_loader):\n",
+    "        print(input_frames.shape, output_frames.shape)  # [B, T, C, H, W]\n",
+    "        break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95e177b0-9b93-42b6-b809-350fadc23a9b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.20"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader-checkpoint.py

@@ -0,0 +1,122 @@
+import numpy as np
+import netCDF4 as nc
+import torch
+import torch.utils.data as data
+
+args = {
+    'data_path': '/data/workspace/yancheng/MM/OriSTP/dataset/05res',
+    'ocean_lead_time': 10,
+    'atmosphere_lead_time': 10,
+    'shuffle': True,
+    'variables_input': [0, 2, 3, 4],
+    'variables_future': [2, 3, 4],
+    'variables_output': [0],
+    'lon_start': 0,
+    'lat_start': 0,
+    'lon_end': 720,
+    'lat_end': 360,
+    'ds_factor': 1,
+}
+
+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 = [(m, n) for m in self.years for n in self.dates]
+        
+    def __getitem__(self, index):
+        years, dates = self.indices[index]
+        train_data = nc.Dataset(f'{self.args["data_path"]}/025res_{years}.nc')
+        input_now = train_data.variables['mhws_variables'][dates-self.args['atmosphere_lead_time']+1:dates+1, 
+                                   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']]
+        
+        input_future = train_data.variables['mhws_variables'][dates+1:dates+self.args['atmosphere_lead_time']+1, 
+                                   self.args['variables_future'],
+                                   self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+                                   self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']]
+        
+        input = np.concatenate([input_now, input_future], 1)
+        
+        target = train_data.variables['mhws_variables'][dates+1:dates+self.args['ocean_lead_time']+1, 
+                                 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']]
+        
+        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)
+        
+        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 = [(m, n) for m in self.years for n in self.dates]
+        
+    def __getitem__(self, index):
+        years, dates = self.indices[index]
+        train_data = nc.Dataset(f'{self.args["data_path"]}/025res_{years}.nc')
+        input_now = train_data.variables['mhws_variables'][dates-self.args['atmosphere_lead_time']+1:dates+1, 
+                                   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']]
+        
+        input_future = train_data.variables['mhws_variables'][dates+1:dates+self.args['atmosphere_lead_time']+1, 
+                                   self.args['variables_future'],
+                                   self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+                                   self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']]
+        
+        input = np.concatenate([input_now, input_future], 1)
+        
+        target = train_data.variables['mhws_variables'][dates+1:dates+self.args['ocean_lead_time']+1, 
+                                 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']]
+        
+        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)
+        
+        return input, target
+
+    def __len__(self):
+        return len(self.indices)
+
+if __name__ == '__main__':
+    args = {
+    'data_path': '/jizhicfs/easyluwu/dataset/ft_local',
+    'ocean_lead_time': 10,
+    'atmosphere_lead_time': 10,
+    'shuffle': True,
+    'variables_input': [1, 2, 3, 4],
+    'variables_future': [2, 3, 4],
+    'variables_output': [1],
+    '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=2)
+    test_loader = data.DataLoader(test_dataset, batch_size=2)
+
+    for inputs, targets in iter(train_loader):
+        print(inputs.shape, targets.shape)
+        break

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_high_kuro-checkpoint.py

@@ -0,0 +1,82 @@
+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 = [(m, n) for m in self.years for n in self.dates]
+        
+    def __getitem__(self, index):
+        years, dates = self.indices[index]
+        train_data = nc.Dataset(f'{self.args["data_path"]}/KURO_{years}_norm.nc')
+        input_now = train_data.variables['mhw_variables'][dates-self.args['atmosphere_lead_time']+1:dates+1, 
+                                   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']]
+        
+        input_future = train_data.variables['mhw_variables'][dates+1:dates+self.args['atmosphere_lead_time']+1, 
+                                   self.args['variables_future'],
+                                   self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+                                   self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']]
+        
+        input = np.concatenate([input_now, input_future], 1)
+        
+        target = train_data.variables['mhw_variables'][dates+1:dates+self.args['ocean_lead_time']+1, 
+                                 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']]
+        
+        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)
+        
+        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, 2021)
+        self.dates = range(12, 357, 3)
+        self.indices = [(m, n) for m in self.years for n in self.dates]
+        
+    def __getitem__(self, index):
+        years, dates = self.indices[index]
+        train_data = nc.Dataset(f'{self.args["data_path"]}/KURO_{years}_norm.nc')
+        input_now = train_data.variables['mhw_variables'][dates-self.args['atmosphere_lead_time']+1:dates+1, 
+                                   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']]
+        
+        input_future = train_data.variables['mhw_variables'][dates+1:dates+self.args['atmosphere_lead_time']+1, 
+                                   self.args['variables_future'],
+                                   self.args['lat_start']:self.args['lat_end']:self.args['ds_factor'],
+                                   self.args['lon_start']:self.args['lon_end']:self.args['ds_factor']]
+        
+        input = np.concatenate([input_now, input_future], 1)
+        
+        target = train_data.variables['mhw_variables'][dates+1:dates+self.args['ocean_lead_time']+1, 
+                                 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']]
+        
+        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)
+        
+        return input, target
+
+    def __len__(self):
+        return len(self.indices)

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio-checkpoint.ipynb

@@ -0,0 +1,209 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "f7a16b9b-07cb-46af-b891-d225ca8a8b2c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/miniconda3/envs/haowu/lib/python3.10/site-packages/torch/cuda/__init__.py:129: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 11000). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /pytorch/c10/cuda/CUDAFunctions.cpp:109.)\n",
+      "  return torch._C._cuda_getDeviceCount() > 0\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])  torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])\n",
+      "\n",
+      "  torch.Size([10, 2, 256, 256])   \n",
+      "torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])\n",
+      "\n",
+      "\n",
+      "\n",
+      " torch.Size([10, 2, 256, 256])\n",
+      "torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])  torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])\n",
+      "\n",
+      "torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256]) \n",
+      "torch.Size([10, 2, 256, 256]) torch.Size([10, 2, 256, 256])  torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])torch.Size([10, 2, 256, 256])\n",
+      "\n",
+      "\n",
+      " torch.Size([10, 2, 256, 256])\n",
+      "torch.Size([10, 2, 256, 256]) torch.Size([10, 2, 256, 256])\n",
+      "torch.Size([1, 10, 2, 256, 256]) torch.Size([1, 10, 2, 256, 256])\n",
+      "输入数据范围: [-1.54, 1.66]\n",
+      "NaN值存在性: False\n",
+      "Inf值存在性: False\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.distributed as dist\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torch.utils.data.distributed import DistributedSampler\n",
+    "import netCDF4 as nc\n",
+    "import numpy as np\n",
+    "\n",
+    "class OceanCurrentDataset(Dataset):\n",
+    "    def __init__(self, data_path, input_steps=10, output_steps=10, transform=None):\n",
+    "        \"\"\"\n",
+    "        海洋流数据集类\n",
+    "        :param data_path: NetCDF文件路径\n",
+    "        :param input_steps: 输入时间步数\n",
+    "        :param output_steps: 预测时间步数\n",
+    "        :param transform: 数据增强变换\n",
+    "        \"\"\"\n",
+    "        self.data_path = data_path\n",
+    "        self.input_steps = input_steps\n",
+    "        self.output_steps = output_steps\n",
+    "        self.transform = transform\n",
+    "        self.total_steps = input_steps + output_steps\n",
+    "        \n",
+    "        # 加载并预处理数据\n",
+    "        self.data = self._load_and_process_data()\n",
+    "        self.mean, self.std = 0, 1\n",
+    "\n",
+    "    def _load_and_process_data(self):\n",
+    "        \"\"\"加载和处理NetCDF数据\"\"\"\n",
+    "        with nc.Dataset(self.data_path, 'r') as ds:\n",
+    "            # 处理缺失值\n",
+    "            def process_var(var):\n",
+    "                arr = var[:]\n",
+    "                if '_FillValue' in var.ncattrs():\n",
+    "                    fill_value = var._FillValue\n",
+    "                    arr = np.ma.masked_values(arr, fill_value).filled(np.nan)\n",
+    "                return torch.nan_to_num(torch.FloatTensor(arr), nan=0.0)\n",
+    "\n",
+    "            # 加载并合并UV分量\n",
+    "            ugos = process_var(ds['ugos'])  # (time, lat, lon)\n",
+    "            vgos = process_var(ds['vgos'])\n",
+    "            \n",
+    "            # 调整维度顺序 [time, channels, lat, lon]\n",
+    "            return torch.stack([ugos, vgos], dim=1)  \n",
+    "\n",
+    "    def _compute_stats(self):\n",
+    "        \"\"\"计算训练集的统计量\"\"\"\n",
+    "        return torch.mean(self.data[:10000]), torch.std(self.data[:10000])\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.data) - self.total_steps + 1\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        window = self.data[idx:idx+self.total_steps]  # [T_total, C, H, W]\n",
+    "        \n",
+    "        window = (window - self.mean) / self.std\n",
+    "        \n",
+    "        # 分割输入输出\n",
+    "        input_seq = window[:self.input_steps]\n",
+    "        target_seq = window[self.input_steps:]\n",
+    "        print(input_seq.shape, target_seq.shape)\n",
+    "        \n",
+    "        if self.transform:\n",
+    "            input_seq = self.transform(input_seq)\n",
+    "            target_seq = self.transform(target_seq)\n",
+    "            \n",
+    "        return input_seq[:,:,::2,::2], target_seq[:,:,::2,::2]\n",
+    "\n",
+    "def create_dataloaders(config):\n",
+    "    full_dataset = OceanCurrentDataset(\n",
+    "        data_path=config['data_path'],\n",
+    "        input_steps=config['input_steps'],\n",
+    "        output_steps=config['output_steps']\n",
+    "    )\n",
+    "    \n",
+    "    train_size = 10000 - config['input_steps'] - config['output_steps'] + 1\n",
+    "    val_size = 500\n",
+    "    test_size = len(full_dataset) - train_size - val_size\n",
+    "    \n",
+    "    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(\n",
+    "        full_dataset, [train_size, val_size, test_size],\n",
+    "        generator=torch.Generator().manual_seed(config['seed'])\n",
+    "    )\n",
+    "    \n",
+    "    # train_sampler = DistributedSampler(train_dataset, shuffle=True)\n",
+    "    # val_sampler = DistributedSampler(val_dataset, shuffle=False)\n",
+    "    # test_sampler = DistributedSampler(test_dataset, shuffle=False)\n",
+    "    \n",
+    "    dataloader_train = DataLoader(\n",
+    "        train_dataset,\n",
+    "        batch_size=config['batch_size'],\n",
+    "        num_workers=config['num_workers'],\n",
+    "        pin_memory=True,\n",
+    "        drop_last=True\n",
+    "    )\n",
+    "    \n",
+    "    dataloader_val = DataLoader(\n",
+    "        val_dataset,\n",
+    "        batch_size=config['val_batch_size'],\n",
+    "        num_workers=config['num_workers'],\n",
+    "        pin_memory=True,\n",
+    "        drop_last=True\n",
+    "    )\n",
+    "    \n",
+    "    dataloader_test = DataLoader(\n",
+    "        test_dataset,\n",
+    "        batch_size=config['val_batch_size'],\n",
+    "        num_workers=config['num_workers'],\n",
+    "        pin_memory=True,\n",
+    "        drop_last=True\n",
+    "    )\n",
+    "    \n",
+    "    return dataloader_train, dataloader_val, dataloader_test, full_dataset.mean, full_dataset.std\n",
+    "\n",
+    "config = {\n",
+    "    'data_path': '/jizhicfs/easyluwu/ocean_project/kuro/KURO.nc',\n",
+    "    'input_steps': 10,\n",
+    "    'output_steps': 10,\n",
+    "    'batch_size': 1,\n",
+    "    'val_batch_size': 1,\n",
+    "    'num_workers': 8,\n",
+    "    'seed': 42\n",
+    "}\n",
+    "# dist.init_process_group(backend='nccl')\n",
+    "\n",
+    "train_loader, val_loader, test_loader, data_mean, data_std = create_dataloaders(config)\n",
+    "\n",
+    "for sample_input, sample_target in train_loader:\n",
+    "    print(sample_input.shape, sample_target.shape)\n",
+    "    print(f\"输入数据范围: [{sample_input.min():.2f}, {sample_input.max():.2f}]\")\n",
+    "    print(f\"NaN值存在性: {torch.isnan(sample_input).any().item()}\")\n",
+    "    print(f\"Inf值存在性: {torch.isinf(sample_input).any().item()}\")\n",
+    "    break"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad0379fc-13ba-48b3-84ad-5356f0e03968",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio-checkpoint.py

@@ -0,0 +1,69 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import h5py
+import numpy as np
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+import torch.utils.data as data
+import h5py
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+class WeatherDataset(Dataset):
+    def __init__(self, data_path, horizon, transform=None):
+        with h5py.File(data_path, 'r') as f:
+            self.data_uv_g = f['uv_g'][:]  
+            self.data_uv_g = torch.from_numpy(self.data_uv_g).to(torch.float32)
+            self.data_uv_g = self.data_uv_g.permute(0, 3, 1, 2).unsqueeze_(2) 
+            
+            self.data_uv_k = f['uv_k'][:]  
+            self.data_uv_k = torch.from_numpy(self.data_uv_k).to(torch.float32)
+            self.data_uv_k = self.data_uv_k.permute(0, 3, 1, 2).unsqueeze_(2) 
+            self.data_uv_gk = torch.cat([self.data_uv_g, self.data_uv_k], dim=2)
+            self.transform = transform
+            self.horizon = horizon
+            self.mean = 0
+            self.std = 1
+    
+    def __len__(self):
+        return len(self.data_uv_gk)
+
+    def __getitem__(self, idx):
+        input_frames = self.data_uv_gk[idx][:self.horizon]
+        output_frames = self.data_uv_gk[idx][self.horizon:2*self.horizon]
+        input_frames = (input_frames - self.mean) / self.std
+        output_frames = (output_frames - self.mean) / self.std
+        return input_frames, output_frames
+
+def load_data(data_path, batch_size, val_batch_size, horizon, num_workers):
+    dataset = WeatherDataset(data_path=data_path+'/kg_all_20_mask_latmean.h5', horizon=horizon, transform=None)
+    dataset_size = len(dataset)
+    train_size = int(dataset_size * 0.8)
+    val_size = int(dataset_size * 0.1)
+    test_size = dataset_size - train_size - val_size
+
+    train_dataset, val_dataset, test_dataset = data.random_split(dataset, [train_size, val_size, test_size])
+
+    train_sampler = DistributedSampler(train_dataset)
+    val_sampler = DistributedSampler(val_dataset)
+    test_sampler = DistributedSampler(test_dataset)
+
+    dataloader_train = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, pin_memory=True,
+                                  num_workers=num_workers, drop_last=True)
+    dataloader_validation = DataLoader(val_dataset, batch_size=val_batch_size, sampler=val_sampler, pin_memory=True,
+                                       num_workers=num_workers, drop_last=True)
+    dataloader_test = DataLoader(test_dataset, batch_size=val_batch_size, sampler=test_sampler, pin_memory=True,
+                                 num_workers=num_workers, drop_last=True)
+    mean, std = 0, 1
+
+    return dataloader_train, dataloader_validation, dataloader_test, mean, std
+
+
+
+
+

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_G_uv-checkpoint.py

@@ -0,0 +1,69 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import h5py
+import numpy as np
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+import torch.utils.data as data
+import h5py
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+class WeatherDataset(Dataset):
+    def __init__(self, data_path, horizon, transform=None):
+        with h5py.File(data_path, 'r') as f:
+            self.data_uv_g = f['u_g'][:]  
+            self.data_uv_g = torch.from_numpy(self.data_uv_g).to(torch.float32)
+            self.data_uv_g = self.data_uv_g.permute(0, 3, 1, 2).unsqueeze_(2) 
+            
+            self.data_uv_k = f['v_g'][:]  
+            self.data_uv_k = torch.from_numpy(self.data_uv_k).to(torch.float32)
+            self.data_uv_k = self.data_uv_k.permute(0, 3, 1, 2).unsqueeze_(2) 
+            self.data_uv_gk = torch.cat([self.data_uv_g, self.data_uv_k], dim=2)
+            self.transform = transform
+            self.horizon = horizon
+            self.mean = 0
+            self.std = 1
+    
+    def __len__(self):
+        return len(self.data_uv_gk)
+
+    def __getitem__(self, idx):
+        input_frames = self.data_uv_gk[idx][:self.horizon]
+        output_frames = self.data_uv_gk[idx][self.horizon:2*self.horizon]
+        input_frames = (input_frames - self.mean) / self.std
+        output_frames = (output_frames - self.mean) / self.std
+        return input_frames, output_frames
+
+def load_data(data_path, batch_size, val_batch_size, horizon, num_workers):
+    dataset = WeatherDataset(data_path=data_path+'/kg_all_20_mask_latmean.h5', horizon=horizon, transform=None)
+    dataset_size = len(dataset)
+    train_size = int(dataset_size * 0.8)
+    val_size = int(dataset_size * 0.1)
+    test_size = dataset_size - train_size - val_size
+
+    train_dataset, val_dataset, test_dataset = data.random_split(dataset, [train_size, val_size, test_size])
+
+    train_sampler = DistributedSampler(train_dataset)
+    val_sampler = DistributedSampler(val_dataset)
+    test_sampler = DistributedSampler(test_dataset)
+
+    dataloader_train = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, pin_memory=False,
+                                  num_workers=num_workers, drop_last=True)
+    dataloader_validation = DataLoader(val_dataset, batch_size=val_batch_size, sampler=val_sampler, pin_memory=False,
+                                       num_workers=num_workers, drop_last=True)
+    dataloader_test = DataLoader(test_dataset, batch_size=val_batch_size, sampler=test_sampler, pin_memory=False,
+                                 num_workers=num_workers, drop_last=True)
+    mean, std = 0, 1
+
+    return dataloader_train, dataloader_validation, dataloader_test, mean, std
+
+
+
+
+

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_K_uv-checkpoint.py

@@ -0,0 +1,69 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import h5py
+import numpy as np
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+import torchvision.transforms as transforms
+import torch.utils.data as data
+import h5py
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+class WeatherDataset(Dataset):
+    def __init__(self, data_path, horizon, transform=None):
+        with h5py.File(data_path, 'r') as f:
+            self.data_uv_g = f['u_k'][:]  
+            self.data_uv_g = torch.from_numpy(self.data_uv_g).to(torch.float32)
+            self.data_uv_g = self.data_uv_g.permute(0, 3, 1, 2).unsqueeze_(2) 
+            
+            self.data_uv_k = f['v_k'][:]  
+            self.data_uv_k = torch.from_numpy(self.data_uv_k).to(torch.float32)
+            self.data_uv_k = self.data_uv_k.permute(0, 3, 1, 2).unsqueeze_(2) 
+            self.data_uv_gk = torch.cat([self.data_uv_g, self.data_uv_k], dim=2)
+            self.transform = transform
+            self.horizon = horizon
+            self.mean = 0
+            self.std = 1
+    
+    def __len__(self):
+        return len(self.data_uv_gk)
+
+    def __getitem__(self, idx):
+        input_frames = self.data_uv_gk[idx][:self.horizon]
+        output_frames = self.data_uv_gk[idx][self.horizon:2*self.horizon]
+        input_frames = (input_frames - self.mean) / self.std
+        output_frames = (output_frames - self.mean) / self.std
+        return input_frames, output_frames
+
+def load_data(data_path, batch_size, val_batch_size, horizon, num_workers):
+    dataset = WeatherDataset(data_path=data_path+'/kg_all_20_mask_latmean.h5', horizon=horizon, transform=None)
+    dataset_size = len(dataset)
+    train_size = int(dataset_size * 0.8)
+    val_size = int(dataset_size * 0.1)
+    test_size = dataset_size - train_size - val_size
+
+    train_dataset, val_dataset, test_dataset = data.random_split(dataset, [train_size, val_size, test_size])
+
+    train_sampler = DistributedSampler(train_dataset)
+    val_sampler = DistributedSampler(val_dataset)
+    test_sampler = DistributedSampler(test_dataset)
+
+    dataloader_train = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, pin_memory=False,
+                                  num_workers=num_workers, drop_last=True)
+    dataloader_validation = DataLoader(val_dataset, batch_size=val_batch_size, sampler=val_sampler, pin_memory=False,
+                                       num_workers=num_workers, drop_last=True)
+    dataloader_test = DataLoader(test_dataset, batch_size=val_batch_size, sampler=test_sampler, pin_memory=False,
+                                 num_workers=num_workers, drop_last=True)
+    mean, std = 0, 1
+
+    return dataloader_train, dataloader_validation, dataloader_test, mean, std
+
+
+
+
+

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_ruiqi-checkpoint.py

@@ -0,0 +1,134 @@
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import netCDF4 as nc
+import numpy as np
+
+class OceanCurrentDataset(Dataset):
+    def __init__(self, data_path, input_steps=10, output_steps=10, transform=None):
+        """
+        海洋流数据集类
+        :param data_path: NetCDF文件路径
+        :param input_steps: 输入时间步数
+        :param output_steps: 预测时间步数
+        :param transform: 数据增强变换
+        """
+        self.data_path = data_path
+        self.input_steps = input_steps
+        self.output_steps = output_steps
+        self.transform = transform
+        self.total_steps = input_steps + output_steps
+        
+        # 加载并预处理数据
+        self.data = self._load_and_process_data()
+        self.mean, self.std = 0, 1
+
+    def _load_and_process_data(self):
+        """加载和处理NetCDF数据"""
+        with nc.Dataset(self.data_path, 'r') as ds:
+            # 处理缺失值
+            def process_var(var):
+                arr = var[:]
+                if '_FillValue' in var.ncattrs():
+                    fill_value = var._FillValue
+                    arr = np.ma.masked_values(arr, fill_value).filled(np.nan)
+                return torch.nan_to_num(torch.FloatTensor(arr), nan=0.0)
+
+            # 加载并合并UV分量
+            ugos = process_var(ds['ugos'])  # (time, lat, lon)
+            vgos = process_var(ds['vgos'])
+            
+            # 调整维度顺序 [time, channels, lat, lon]
+            return torch.stack([ugos, vgos], dim=1)  
+
+    def _compute_stats(self):
+        """计算训练集的统计量"""
+        return torch.mean(self.data[:10000]), torch.std(self.data[:10000])
+
+    def __len__(self):
+        return len(self.data) - self.total_steps + 1
+
+    def __getitem__(self, idx):
+        window = self.data[idx:idx+self.total_steps]  # [T_total, C, H, W]
+        
+        window = (window - self.mean) / self.std
+        
+        # 分割输入输出
+        input_seq = window[:self.input_steps]
+        target_seq = window[self.input_steps:]
+        
+        if self.transform:
+            input_seq = self.transform(input_seq)
+            target_seq = self.transform(target_seq)
+            
+        return input_seq, target_seq
+
+def create_dataloaders(config):
+    full_dataset = OceanCurrentDataset(
+        data_path=config['data_path'],
+        input_steps=config['input_steps'],
+        output_steps=config['output_steps']
+    )
+    
+    train_size = 10000 - config['input_steps'] - config['output_steps'] + 1
+    val_size = 500
+    test_size = len(full_dataset) - train_size - val_size
+    
+    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
+        full_dataset, [train_size, val_size, test_size],
+        generator=torch.Generator().manual_seed(config['seed'])
+    )
+    
+    train_sampler = DistributedSampler(train_dataset, shuffle=True)
+    val_sampler = DistributedSampler(val_dataset, shuffle=False)
+    test_sampler = DistributedSampler(test_dataset, shuffle=False)
+    
+    dataloader_train = DataLoader(
+        train_dataset,
+        batch_size=config['batch_size'],
+        sampler=train_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_val = DataLoader(
+        val_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=val_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_test = DataLoader(
+        test_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=test_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    return dataloader_train, dataloader_val, dataloader_test, full_dataset.mean, full_dataset.std
+
+# config = {
+#     'data_path': '/jizhicfs/easyluwu/ocean_project/kuro/KURO.nc',
+#     'input_steps': 10,
+#     'output_steps': 10,
+#     'batch_size': 1,
+#     'val_batch_size': 1,
+#     'num_workers': 8,
+#     'seed': 42
+# }
+# dist.init_process_group(backend='nccl')
+
+# train_loader, val_loader, test_loader, data_mean, data_std = create_dataloaders(config)
+
+# for sample_input, sample_target in train_loader:
+#     print(sample_input.shape, sample_target.shape)
+#     print(f"输入数据范围: [{sample_input.min():.2f}, {sample_input.max():.2f}]")
+#     print(f"NaN值存在性: {torch.isnan(sample_input).any().item()}")
+#     print(f"Inf值存在性: {torch.isinf(sample_input).any().item()}")
+#     break

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_ruiqi_128-checkpoint.py

@@ -0,0 +1,134 @@
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import netCDF4 as nc
+import numpy as np
+
+class OceanCurrentDataset(Dataset):
+    def __init__(self, data_path, input_steps=10, output_steps=10, transform=None):
+        """
+        海洋流数据集类
+        :param data_path: NetCDF文件路径
+        :param input_steps: 输入时间步数
+        :param output_steps: 预测时间步数
+        :param transform: 数据增强变换
+        """
+        self.data_path = data_path
+        self.input_steps = input_steps
+        self.output_steps = output_steps
+        self.transform = transform
+        self.total_steps = input_steps + output_steps
+        
+        # 加载并预处理数据
+        self.data = self._load_and_process_data()
+        self.mean, self.std = 0, 1
+
+    def _load_and_process_data(self):
+        """加载和处理NetCDF数据"""
+        with nc.Dataset(self.data_path, 'r') as ds:
+            # 处理缺失值
+            def process_var(var):
+                arr = var[:]
+                if '_FillValue' in var.ncattrs():
+                    fill_value = var._FillValue
+                    arr = np.ma.masked_values(arr, fill_value).filled(np.nan)
+                return torch.nan_to_num(torch.FloatTensor(arr), nan=0.0)
+
+            # 加载并合并UV分量
+            ugos = process_var(ds['ugos'])  # (time, lat, lon)
+            vgos = process_var(ds['vgos'])
+            
+            # 调整维度顺序 [time, channels, lat, lon]
+            return torch.stack([ugos, vgos], dim=1)  
+
+    def _compute_stats(self):
+        """计算训练集的统计量"""
+        return torch.mean(self.data[:10000]), torch.std(self.data[:10000])
+
+    def __len__(self):
+        return len(self.data) - self.total_steps + 1
+
+    def __getitem__(self, idx):
+        window = self.data[idx:idx+self.total_steps]  # [T_total, C, H, W]
+        
+        window = (window - self.mean) / self.std
+        
+        # 分割输入输出
+        input_seq = window[:self.input_steps]
+        target_seq = window[self.input_steps:]
+        
+        if self.transform:
+            input_seq = self.transform(input_seq)
+            target_seq = self.transform(target_seq)
+            
+        return input_seq[:,:,::2,::2], target_seq[:,:,::2,::2]
+
+def create_dataloaders(config):
+    full_dataset = OceanCurrentDataset(
+        data_path=config['data_path'],
+        input_steps=config['input_steps'],
+        output_steps=config['output_steps']
+    )
+    
+    train_size = 10000 - config['input_steps'] - config['output_steps'] + 1
+    val_size = 500
+    test_size = len(full_dataset) - train_size - val_size
+    
+    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
+        full_dataset, [train_size, val_size, test_size],
+        generator=torch.Generator().manual_seed(config['seed'])
+    )
+    
+    train_sampler = DistributedSampler(train_dataset, shuffle=True)
+    val_sampler = DistributedSampler(val_dataset, shuffle=False)
+    test_sampler = DistributedSampler(test_dataset, shuffle=False)
+    
+    dataloader_train = DataLoader(
+        train_dataset,
+        batch_size=config['batch_size'],
+        sampler=train_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_val = DataLoader(
+        val_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=val_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_test = DataLoader(
+        test_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=test_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    return dataloader_train, dataloader_val, dataloader_test, full_dataset.mean, full_dataset.std
+
+# config = {
+#     'data_path': '/jizhicfs/easyluwu/ocean_project/kuro/KURO.nc',
+#     'input_steps': 10,
+#     'output_steps': 10,
+#     'batch_size': 1,
+#     'val_batch_size': 1,
+#     'num_workers': 8,
+#     'seed': 42
+# }
+# dist.init_process_group(backend='nccl')
+
+# train_loader, val_loader, test_loader, data_mean, data_std = create_dataloaders(config)
+
+# for sample_input, sample_target in train_loader:
+#     print(sample_input.shape, sample_target.shape)
+#     print(f"输入数据范围: [{sample_input.min():.2f}, {sample_input.max():.2f}]")
+#     print(f"NaN值存在性: {torch.isnan(sample_input).any().item()}")
+#     print(f"Inf值存在性: {torch.isinf(sample_input).any().item()}")
+#     break

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_ruiqi_64-checkpoint.py

@@ -0,0 +1,134 @@
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import netCDF4 as nc
+import numpy as np
+
+class OceanCurrentDataset(Dataset):
+    def __init__(self, data_path, input_steps=10, output_steps=10, transform=None):
+        """
+        海洋流数据集类
+        :param data_path: NetCDF文件路径
+        :param input_steps: 输入时间步数
+        :param output_steps: 预测时间步数
+        :param transform: 数据增强变换
+        """
+        self.data_path = data_path
+        self.input_steps = input_steps
+        self.output_steps = output_steps
+        self.transform = transform
+        self.total_steps = input_steps + output_steps
+        
+        # 加载并预处理数据
+        self.data = self._load_and_process_data()
+        self.mean, self.std = 0, 1
+
+    def _load_and_process_data(self):
+        """加载和处理NetCDF数据"""
+        with nc.Dataset(self.data_path, 'r') as ds:
+            # 处理缺失值
+            def process_var(var):
+                arr = var[:]
+                if '_FillValue' in var.ncattrs():
+                    fill_value = var._FillValue
+                    arr = np.ma.masked_values(arr, fill_value).filled(np.nan)
+                return torch.nan_to_num(torch.FloatTensor(arr), nan=0.0)
+
+            # 加载并合并UV分量
+            ugos = process_var(ds['ugos'])  # (time, lat, lon)
+            vgos = process_var(ds['vgos'])
+            
+            # 调整维度顺序 [time, channels, lat, lon]
+            return torch.stack([ugos, vgos], dim=1)  
+
+    def _compute_stats(self):
+        """计算训练集的统计量"""
+        return torch.mean(self.data[:10000]), torch.std(self.data[:10000])
+
+    def __len__(self):
+        return len(self.data) - self.total_steps + 1
+
+    def __getitem__(self, idx):
+        window = self.data[idx:idx+self.total_steps]  # [T_total, C, H, W]
+        
+        window = (window - self.mean) / self.std
+        
+        # 分割输入输出
+        input_seq = window[:self.input_steps]
+        target_seq = window[self.input_steps:]
+        
+        if self.transform:
+            input_seq = self.transform(input_seq)
+            target_seq = self.transform(target_seq)
+            
+        return input_seq[:,:,::4,::4], target_seq[:,:,::4,::4]
+
+def create_dataloaders(config):
+    full_dataset = OceanCurrentDataset(
+        data_path=config['data_path'],
+        input_steps=config['input_steps'],
+        output_steps=config['output_steps']
+    )
+    
+    train_size = 10000 - config['input_steps'] - config['output_steps'] + 1
+    val_size = 500
+    test_size = len(full_dataset) - train_size - val_size
+    
+    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
+        full_dataset, [train_size, val_size, test_size],
+        generator=torch.Generator().manual_seed(config['seed'])
+    )
+    
+    train_sampler = DistributedSampler(train_dataset, shuffle=True)
+    val_sampler = DistributedSampler(val_dataset, shuffle=False)
+    test_sampler = DistributedSampler(test_dataset, shuffle=False)
+    
+    dataloader_train = DataLoader(
+        train_dataset,
+        batch_size=config['batch_size'],
+        sampler=train_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_val = DataLoader(
+        val_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=val_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_test = DataLoader(
+        test_dataset,
+        batch_size=config['val_batch_size'],
+        sampler=test_sampler,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    return dataloader_train, dataloader_val, dataloader_test, full_dataset.mean, full_dataset.std
+
+# config = {
+#     'data_path': '/jizhicfs/easyluwu/ocean_project/kuro/KURO.nc',
+#     'input_steps': 10,
+#     'output_steps': 10,
+#     'batch_size': 1,
+#     'val_batch_size': 1,
+#     'num_workers': 8,
+#     'seed': 42
+# }
+# dist.init_process_group(backend='nccl')
+
+# train_loader, val_loader, test_loader, data_mean, data_std = create_dataloaders(config)
+
+# for sample_input, sample_target in train_loader:
+#     print(sample_input.shape, sample_target.shape)
+#     print(f"输入数据范围: [{sample_input.min():.2f}, {sample_input.max():.2f}]")
+#     print(f"NaN值存在性: {torch.isnan(sample_input).any().item()}")
+#     print(f"Inf值存在性: {torch.isinf(sample_input).any().item()}")
+#     break

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_kuroshio_ruiqi_single-checkpoint.py

@@ -0,0 +1,96 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import netCDF4 as nc
+import numpy as np
+
+class OceanCurrentDataset(Dataset):
+    def __init__(self, data_path, input_steps=10, output_steps=10, transform=None):
+        self.data_path = data_path
+        self.input_steps = input_steps
+        self.output_steps = output_steps
+        self.transform = transform
+        self.total_steps = input_steps + output_steps
+        
+        # Load and process data
+        self.data = self._load_and_process_data()
+        self.mean, self.std = self._compute_stats()
+
+    def _load_and_process_data(self):
+        """Load and process NetCDF data"""
+        with nc.Dataset(self.data_path, 'r') as ds:
+            def process_var(var):
+                arr = var[:]
+                if '_FillValue' in var.ncattrs():
+                    fill_value = var._FillValue
+                    arr = np.ma.masked_values(arr, fill_value).filled(np.nan)
+                return torch.nan_to_num(torch.FloatTensor(arr), nan=0.0)
+
+            ugos = process_var(ds['ugos'])  # (time, lat, lon)
+            vgos = process_var(ds['vgos'])
+            
+            return torch.stack([ugos, vgos], dim=1)  # [time, channels, lat, lon]
+
+    def _compute_stats(self):
+        """Compute dataset statistics"""
+        return torch.mean(self.data[:10000]), torch.std(self.data[:10000])
+
+    def __len__(self):
+        return len(self.data) - self.total_steps + 1
+
+    def __getitem__(self, idx):
+        window = self.data[idx:idx+self.total_steps]  # [T_total, C, H, W]
+        window = (window - 0) / 1
+        
+        input_seq = window[:self.input_steps]
+        target_seq = window[self.input_steps:]
+        
+        if self.transform:
+            input_seq = self.transform(input_seq)
+            target_seq = self.transform(target_seq)
+            
+        return input_seq, target_seq
+
+def create_dataloaders(config):
+    full_dataset = OceanCurrentDataset(
+        data_path=config['data_path'],
+        input_steps=config['input_steps'],
+        output_steps=config['output_steps']
+    )
+    
+    train_size = 10000 - config['input_steps'] - config['output_steps'] + 1
+    val_size = 500
+    test_size = len(full_dataset) - train_size - val_size
+    
+    train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(
+        full_dataset, [train_size, val_size, test_size],
+        generator=torch.Generator().manual_seed(config['seed'])
+    )
+    
+    dataloader_train = DataLoader(
+        train_dataset,
+        batch_size=config['batch_size'],
+        shuffle=True,  # Changed from DistributedSampler
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_val = DataLoader(
+        val_dataset,
+        batch_size=config['val_batch_size'],
+        shuffle=False,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    dataloader_test = DataLoader(
+        test_dataset,
+        batch_size=config['val_batch_size'],
+        shuffle=False,
+        num_workers=config['num_workers'],
+        pin_memory=True,
+        drop_last=True
+    )
+    
+    return dataloader_train, dataloader_val, dataloader_test,0, 1

Exp3_Kuroshio_forecasting/dataloader_api/.ipynb_checkpoints/dataloader_test-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}