Merge pull request #521 from AIStream-Peelout/gpu_class_bug

DSANet + GPU bug

.circleci/config.yml

@@ -289,6 +289,7 @@ jobs:
             coverage run flood_forecast/trainer.py -p tests/custom_encode.json
             coverage run flood_forecast/trainer.py -p tests/multi_decoder_test.json
             coverage run flood_forecast/trainer.py -p tests/test_dual.json
+            coverage run flood_forecast/trainer.py -p tests/dsanet.json
       - store_test_results:
           path: test-results
 

README.md

@@ -29,6 +29,7 @@ Using the library
 7. [Transformer XL](https://arxiv.org/abs/1901.02860): Porting Transformer XL for time series.
 8. [Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting](https://arxiv.org/abs/2012.07436) (Informer)
 9. [DeepAR](https://arxiv.org/abs/1704.04110)
+10. [DSANet](https://www.semanticscholar.org/paper/DSANet%3A-Dual-Self-Attention-Network-for-Time-Series-Huang-Wang/6645a09c742760144e4ba0a6f6652e429b1bf107)
 
 **Forthcoming Models**
 

flood_forecast/model_dict_function.py

@@ -17,6 +17,7 @@
 from flood_forecast.transformer_xl.transformer_bottleneck import DecoderTransformer
 from flood_forecast.custom.dilate_loss import DilateLoss
 from flood_forecast.meta_models.basic_ae import AE
+from flood_forecast.transformer_xl.dsanet import DSANet
 
 """
 Utility dictionaries to map a string to a class
@@ -32,7 +33,8 @@
     "DARNN": DARNN,
     "DecoderTransformer": DecoderTransformer,
     "BasicAE": AE,
-    "Informer": Informer
+    "Informer": Informer,
+    "DSANet": DSANet
 }
 
 pytorch_criterion_dict = {

flood_forecast/preprocessing/process_usgs.py

@@ -7,7 +7,7 @@
 
 
 def make_usgs_data(start_date: datetime, end_date: datetime, site_number: str) -> pd.DataFrame:
-    """ """
+    """"""
     base_url = "https://nwis.waterdata.usgs.gov/usa/nwis/uv/?cb_00060=on&cb_00065&format=rdb&"
     full_url = base_url + "site_no=" + site_number + "&period=&begin_date=" + \
         start_date.strftime("%Y-%m-%d") + "&end_date=" + end_date.strftime("%Y-%m-%d")

flood_forecast/pytorch_training.py

@@ -591,7 +591,7 @@ def compute_validation(validation_loader: DataLoader,
         wandb.log({"roc_" + str(epoch): wandb.plot.roc_curve(fin, mod_output1, classes_to_plot=None, labels=None,
                                                              title="roc_" + str(epoch))})
         wandb.log({"pr": wandb.plot.pr_curve(fin, mod_output1)})
-        wandb.log({"conf_": wandb.plot.confusion_matrix(probs=mod_output1.detach().numpy(), y_true=fin.detach().numpy(),
-                   class_names=None)})
+        wandb.log({"conf_": wandb.plot.confusion_matrix(probs=mod_output1.detach().cpu().numpy(),
+                   y_true=fin.detach().cpu().numpy(), class_names=None)})
     model.train()
     return list(scaled_crit.values())[0]

flood_forecast/time_model.py

@@ -172,7 +172,7 @@ def save_model(self, final_path: str, epoch: int) -> None:
                 print("Wandb stupid error")
                 print(e.__traceback__)
 
-    def __re_add_params__(self, start_end_params, dataset_params, data_path):
+    def __re_add_params__(self, start_end_params: Dict, dataset_params, data_path):
         """
         Function to re-add the params to the model
         """