add new channel support

Author: howard-hou

Pretrained-RWKV_TS/src/dataset.py

@@ -25,12 +25,21 @@ def build_test_set(self):
         data_df = df[keys_sorted[-1]]
         X = data_df["nwp_ws100"].to_numpy()[:, np.newaxis]
         y = data_df["fj_windSpeed"].to_numpy()[:, np.newaxis]
+        print(f"input shape: {X.shape}, target shape: {y.shape}")
+        # shift the fj_windSpeed
+        shifted_list = []
+        for i in range(1, self.args.shift_steps+1):
+            shifted_windspeed = data_df["fj_windSpeed"].shift(i).fillna(0).to_numpy()[:, np.newaxis]
+            shifted_list.append(shifted_windspeed)
+        self.shifted_y = np.concatenate(shifted_list, axis=1)
         # split X, y to chunk 
         X_chunks = [X[i-self.prefix_len:i+self.seq_len] for i in range(0, len(X), self.seq_len) if i != 0]
         y_chunks = [y[i:i+self.seq_len] for i in range(0, len(y), self.seq_len) if i != 0]
+        y_shifted_chunks = [self.shifted_y[i:i+self.seq_len] for i in range(0, len(self.shifted_y), self.seq_len) if i != 0]
         print(f"input chunks: {len(X_chunks)}, target chunks: {len(y_chunks)}")
         self.X = X_chunks
         self.y = y_chunks
+        self.shifted_y = y_shifted_chunks
 
     def __len__(self):
         return len(self.X)
@@ -41,7 +50,8 @@ def __getitem__(self, idx):
         else:
             input_points = self.X[idx]
         targets = self.y[idx]
-        return dict(input_points=input_points, targets=targets)
+        shifted_targets = self.shifted_y[idx]
+        return dict(input_points=input_points, targets=targets, shifted_targets=shifted_targets)
 
 
 class TrainDataset(Dataset):
@@ -76,6 +86,12 @@ def build_train_set(self):
         self.X_std = self.X.std()
         self.y_mean = self.y.mean()
         self.y_std = self.y.std()
+        # shift the fj_windSpeed
+        shifted_list = []
+        for i in range(1, self.args.shift_steps+1):
+            shifted_windspeed = data_df["fj_windSpeed"].shift(i).fillna(0).to_numpy()[:, np.newaxis]
+            shifted_list.append(shifted_windspeed)
+        self.shifted_y = np.concatenate(shifted_list, axis=1)
 
     def __len__(self):
         return self.args.epoch_steps * self.args.micro_bsz
@@ -86,7 +102,9 @@ def __getitem__(self, idx):
         if self.do_normalize:
             input_points = (self.X[s-self.prefix_len:s+self.seq_len] - self.X_mean) / self.X_std
             targets = (self.y[s:s+self.seq_len] - self.y_mean) / self.y_std
+            shifted_targets = (self.shifted_y[s:s+self.seq_len] - self.y_mean) / self.y_std
         else:
             input_points = self.X[s-self.prefix_len:s+self.seq_len] # include the previous seq_len points
             targets = self.y[s:s+self.seq_len]
-        return dict(input_points=input_points, targets=targets)
+            shifted_targets = self.shifted_y[s:s+self.seq_len]
+        return dict(input_points=input_points, targets=targets, shifted_targets=shifted_targets)

Pretrained-RWKV_TS/src/model.py

@@ -342,7 +342,26 @@ def __init__(self, n_embd):
         self.drop = nn.Dropout(0.1)
 
     def forward(self, x):
-        return self.drop(torch.cat([self.conv3(x), self.conv5(x), self.conv7(x), self.conv9(x)], 1))
+        B, T, C = x.size()
+        return self.drop(torch.cat([self.conv3(x), self.conv5(x), self.conv7(x), self.conv9(x)], 1)).view(B, T, -1)
+
+
+class Projector2(nn.Module):
+    def __init__(self, n_embd, shift_steps):
+        super().__init__()
+        self.conv3 = nn.Conv1d(1, n_embd//8, kernel_size=3, stride=1, groups=1, padding='same')
+        self.conv5 = nn.Conv1d(1, n_embd//8, kernel_size=5, stride=1, groups=1, padding='same')
+        self.conv7 = nn.Conv1d(1, n_embd//8, kernel_size=7, stride=1, groups=1, padding='same')
+        self.conv9 = nn.Conv1d(1, n_embd//8, kernel_size=9, stride=1, groups=1, padding='same')
+        self.linear = nn.Linear(shift_steps, n_embd//2)
+        self.drop = nn.Dropout(0.1)
+
+    def forward(self, x, shifted_y):
+        B, T, C = x.size()
+        x = x.view(B, C, T)
+        shifted_y_emb = self.linear(shifted_y)
+        x_emb = torch.cat([self.conv3(x), self.conv5(x), self.conv7(x), self.conv9(x)], 1)
+        return self.drop(torch.cat([x_emb.view(B, T, -1), shifted_y_emb], 2))
 
 
 
@@ -353,7 +372,7 @@ def __init__(self, args):
         self.rwkv = RWKV(args)
         if args.load_model:
             self.load_rwkv_from_pretrained(args.load_model)
-        self.proj = Projector(args.n_embd)
+        self.proj = Projector2(args.n_embd, args.shift_steps)
         self.head = nn.Linear(args.n_embd, 1, bias=False)
         self.best_val_loss = torch.tensor(float("inf"))
         self.do_normalize = args.do_normalize
@@ -393,12 +412,11 @@ def configure_optimizers(self):
         return FusedAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adam_w_mode=True, amsgrad=False)
 
     def forward(self, samples):
-        x, targets = samples["input_points"], samples["targets"]
-        B, T, C = x.shape
-        x = self.proj(x.view(B, C, T)).view(B, T, -1)
+        x, y, shifted_y = samples["input_points"], samples["targets"], samples["shifted_targets"]
+        x = self.proj(x, shifted_y)
         x = self.rwkv(x)[:, self.prefix_len:, :] # 
         outputs = self.head(x)
-        return outputs, targets
+        return outputs, y
 
     def bidirectional_forward(self, x, x_emb=None):
         pass

Pretrained-RWKV_TS/train.py

@@ -57,6 +57,7 @@
     parser.add_argument("--label_smoothing", default=0, type=int)  # label smoothing window
     parser.add_argument("--prefix_len", default=0, type=int)  # 
     parser.add_argument("--do_normalize", action="store_true")  # normalize input
+    parser.add_argument("--shift_steps", default=5, type=int)  # shift steps for fj_windSpeed
 
     parser = Trainer.add_argparse_args(parser)
     args = parser.parse_args()