cws

杭州日电量/industry_elec_cws.py

@@ -42,6 +42,8 @@ def to_data(file_dir, excel):
     data.sort_values(by='stat_date', ascending=True)
     data.drop(columns=[i for i in data.columns if (data[i] == 0).sum() / len(data) >= 0.5], inplace=True)  # 去除0值列
     print('len(data):', len(data))
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("运算类型：", device)
     for industry in data.columns[1:]:
         c = time.time()
         df = data[['stat_date', industry]]
@@ -71,22 +73,17 @@ def to_data(file_dir, excel):
         train_x = train_x.reshape(-1, 1, DAYS_FOR_TRAIN)
         train_y = train_y.reshape(-1, 1, 1)
 
-        # 使用GPU
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
         # 转为pytorch的tensor对象
         train_x = torch.from_numpy(train_x)
         train_y = torch.from_numpy(train_y)
 
-        train_x = train_x.cuda()
-        train_y = train_y.cuda()
         model = LSTM_Regression(DAYS_FOR_TRAIN, 32, output_size=1, num_layers=2)  # 导入模型并设置模型的参数输入输出层、隐藏层等
 
-        # 使用GPU
+        # 训练使用GPU
-        print("cuda" if torch.cuda.is_available() else "cpu")
+        if torch.cuda.is_available():
-        model.to(device)
+            train_x = train_x.cuda()
-        # train_x.to(device)
+            train_y = train_y.cuda()
-        # train_y.to(device)
+            model.to(device)
 
         train_loss = []
         loss_function = nn.MSELoss()
@@ -115,11 +112,20 @@ def to_data(file_dir, excel):
 
         dataset_x = dataset_x.reshape(-1, 1, DAYS_FOR_TRAIN)  # (seq_size, batch_size, feature_size)
         dataset_x = torch.from_numpy(dataset_x)
-        dataset_x = dataset_x.cuda()
+
+        # 测试使用GPU
+        if torch.cuda.is_available():
+            dataset_x = dataset_x.cuda()
 
         pred_test = model(dataset_x)  # 全量训练集
         # 模型输出 (seq_size, batch_size, output_size)
-        pred_test = pred_test.cpu()
+
+        # 测试使用GPU ######################################################### 注意调整这里，反复转换效率不高
+        if torch.cuda.is_available():
+            # 不支持将GPU tensor转换为numpy
+            pred_test = pred_test.cpu()
+        # 测试使用GPU ######################################################### 注意调整这里，反复转换效率不高
+
         pred_test = pred_test.view(-1).data.numpy()
         pred_test = np.concatenate((np.zeros(DAYS_FOR_TRAIN), pred_test))
         assert len(pred_test) == len(df)