feat: onnx model [PDS-570]

.gitignore

@@ -4,3 +4,5 @@ __pycache__
 training/*
 scripts/_nohup
 .rsync-filter
+.DS_Store
+launch.json

models/encoder.py

@@ -40,29 +40,33 @@ def __init__(self, input_dims, output_dims, hidden_dims=64, kernel_size=3, depth
 
     def forward(self, x, mask=None):  # x: B x T x input_dims
         nan_mask = ~x.isnan().any(axis=-1)
+        # nan_mask = ~torch.any(torch.isnan(x), axis=-1)
         x[~nan_mask] = 0
         x = self.input_fc(x)  # B x T x Ch
 
         # generate & apply mask
-        if mask is None:
-            if self.training:
-                mask = self.mask_mode
-            else:
-                mask = 'all_true'
+        # if mask is None:
+        #     if self.training:
+        #         mask = self.mask_mode
+        #     else:
+        #         mask = 'all_true'
 
-        if mask == 'binomial':
-            mask = generate_binomial_mask(x.size(0), x.size(1)).to(x.device)
-        elif mask == 'continuous':
-            mask = generate_continuous_mask(x.size(0), x.size(1)).to(x.device)
-        elif mask == 'all_true':
-            mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
-        elif mask == 'all_false':
-            mask = x.new_full((x.size(0), x.size(1)), False, dtype=torch.bool)
-        elif mask == 'mask_last':
-            mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
-            mask[:, -1] = False
+        # if mask == 'binomial':
+        #     mask = generate_binomial_mask(x.size(0), x.size(1)).to(x.device)
+        # elif mask == 'continuous':
+        #     mask = generate_continuous_mask(x.size(0), x.size(1)).to(x.device)
+        # elif mask == 'all_true':
+        #     mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
+        # elif mask == 'all_false':
+        #     mask = x.new_full((x.size(0), x.size(1)), False, dtype=torch.bool)
+        # elif mask == 'mask_last':
+        #     mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
+        #     mask[:, -1] = False
 
-        mask &= nan_mask
+        if mask is None: 
+            mask = nan_mask
+        else:
+            mask &= nan_mask
         x[~mask] = 0
 
         # conv encoder

onnx_test.py

@@ -0,0 +1,81 @@
+import os
+import pickle 
+from datetime import datetime
+
+import numpy as np 
+import pandas as pd 
+# import matplotlib.pyplot as plt
+
+from ts2vec import TS2Vec
+import torch
+from torch import nn
+
+
+# Model traing parameters 
+OUTPUT_DIMS = 320 
+TEMPORAL_UNIT = 2 
+BATCH_SIZE = 128
+N_EPOCHS = 2
+HIDDEN_DIMS = 64
+KERNEL_SIZE = 3
+SAVE_CHECK_POINT = False
+
+
+# load train data 
+train_data_dir = "/Users/fguo/cmt/ts2vec/sample_data/train_data_sample10.npy"
+train_motion_names_dir = "/Users/fguo/cmt/ts2vec/sample_data/train_motion_names_sample10.parquet"
+train_data = np.load(train_data_dir)
+train_motion_names = pd.read_parquet(train_motion_names_dir)
+print(train_data.shape, train_motion_names.shape)
+
+# load val data 
+val_data_dir = "/Users/fguo/cmt/ts2vec/sample_data/val_data_sample4.npy"
+val_motion_names_dir = "/Users/fguo/cmt/ts2vec/sample_data/val_motion_names_sample4.parquet"
+val_data = np.load(val_data_dir)
+val_motion_names = pd.read_parquet(val_motion_names_dir)
+print(val_data.shape, val_motion_names.shape)
+
+
+model = TS2Vec(
+    input_dims=train_data.shape[-1],
+    device='cpu', 
+    output_dims=OUTPUT_DIMS,
+    hidden_dims=HIDDEN_DIMS, 
+    temporal_unit=TEMPORAL_UNIT,
+    batch_size=BATCH_SIZE,
+    after_epoch_callback=None
+)
+# loss_log = model.fit(
+#     train_data,
+#     val_data,
+#     verbose=True, 
+#     n_epochs=N_EPOCHS, 
+#     )
+# model.save("/Users/fguo/cmt/ts2vec/model_checkpoints/sample_model.pkl")
+
+model.load("/Users/fguo/cmt/ts2vec/model_checkpoints/sample_model.pkl")
+random_input = torch.randn(1, 1, 8, device="cpu")
+# out = model.encode(random_input)
+
+
+# class MLPModel(nn.Module):
+#   def __init__(self):
+#       super().__init__()
+#       self.fc0 = nn.Linear(8, 8, bias=True)
+#       self.fc1 = nn.Linear(8, 4, bias=True)
+#       self.fc2 = nn.Linear(4, 2, bias=True)
+#       self.fc3 = nn.Linear(2, 2, bias=True)
+
+#   def forward(self, tensor_x: torch.Tensor):
+#       tensor_x = self.fc0(tensor_x)
+#       tensor_x = torch.sigmoid(tensor_x)
+#       tensor_x = self.fc1(tensor_x)
+#       tensor_x = torch.sigmoid(tensor_x)
+#       tensor_x = self.fc2(tensor_x)
+#       tensor_x = torch.sigmoid(tensor_x)
+#       output = self.fc3(tensor_x)
+#       return output
+
+# model = MLPModel()
+# tensor_x = torch.rand((97, 8), dtype=torch.float32)
+# onnx_program = torch.onnx.dynamo_export(model, tensor_x)

requirements.txt

@@ -1,7 +1,13 @@
-Bottleneck==1.3.2
-torch==1.8.1
-scipy==1.6.1
-numpy==1.19.2
-statsmodels==0.12.2
-pandas==1.0.1
-scikit_learn==0.24.2
+Bottleneck
+torch>=2.0
+scipy>=1.6.1
+numpy>=1.19.2
+statsmodels>=0.12.2
+pandas>=1.0.1
+scikit_learn>=0.24.2
+onnxruntime>=1.12.0
+ipykernel
+pyarrow
+fastparquet
+onnxscript
+onnxoptimizer

ts2vec.py

@@ -6,6 +6,7 @@
 from models.losses import hierarchical_contrastive_loss
 from utils import take_per_row, split_with_nan, centerize_vary_length_series, torch_pad_nan
 import math
+from models.encoder import generate_binomial_mask
 
 class TS2Vec:
     '''The TS2Vec model'''
@@ -124,10 +125,14 @@ def fit(self, train_data, val_data, n_epochs=None, n_iters=None, verbose=False):
 
                 optimizer.zero_grad()
 
-                out1 = self._net(take_per_row(x, crop_offset + crop_eleft, crop_right - crop_eleft))
+                in1 = take_per_row(x, crop_offset + crop_eleft, crop_right - crop_eleft)
+                mask = generate_binomial_mask(in1.size(0), in1.size(1)).to(self.device)
+                out1 = self._net(in1, mask)
                 out1 = out1[:, -crop_l:]
 
-                out2 = self._net(take_per_row(x, crop_offset + crop_left, crop_eright - crop_left))
+                in2 = take_per_row(x, crop_offset + crop_left, crop_eright - crop_left)
+                mask = generate_binomial_mask(in2.size(0), in2.size(1)).to(self.device)
+                out2 = self._net(in2, mask)
                 out2 = out2[:, :crop_l]
 
                 loss = hierarchical_contrastive_loss(
@@ -204,11 +209,15 @@ def evaluate(self, val_data):
                 crop_eleft = np.random.randint(crop_left + 1)
                 crop_eright = np.random.randint(low=crop_right, high=ts_l + 1)
                 crop_offset = np.random.randint(low=-crop_eleft, high=ts_l - crop_eright + 1, size=x.size(0))
-
-                out1 = self._net(take_per_row(x, crop_offset + crop_eleft, crop_right - crop_eleft))
+
+                in1 = take_per_row(x, crop_offset + crop_eleft, crop_right - crop_eleft)
+                mask = generate_binomial_mask(in1.size(0), in1.size(1)).to(self.device)
+                out1 = self._net(in1, mask)
                 out1 = out1[:, -crop_l:]
-
-                out2 = self._net(take_per_row(x, crop_offset + crop_left, crop_eright - crop_left))
+
+                in2 = take_per_row(x, crop_offset + crop_left, crop_eright - crop_left)
+                mask = generate_binomial_mask(in2.size(0), in2.size(1)).to(self.device)
+                out2 = self._net(in2, mask)
                 out2 = out2[:, :crop_l]
 
                 loss_val = hierarchical_contrastive_loss(