feat: use gpt2 for base model [PDS-512]

.gitignore

@@ -6,3 +6,6 @@ datasets
 *.log
 *.pt
 *.zip
+*.npy 
+*.pt
+*.parquet

code/TFC/embed.py

@@ -0,0 +1,185 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+import math
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEmbedding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float()
+                    * -(math.log(10000.0) / d_model)).exp()
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe[:, :x.size(1)]
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(TokenEmbedding, self).__init__()
+        padding = 1 if torch.__version__ >= '1.5.0' else 2
+        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
+                                   kernel_size=3, padding=padding, padding_mode='circular', bias=False)
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_in', nonlinearity='leaky_relu')
+
+    def forward(self, x):
+        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
+        return x
+
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(FixedEmbedding, self).__init__()
+
+        w = torch.zeros(c_in, d_model).float()
+        w.require_grad = False
+
+        position = torch.arange(0, c_in).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float()
+                    * -(math.log(10000.0) / d_model)).exp()
+
+        w[:, 0::2] = torch.sin(position * div_term)
+        w[:, 1::2] = torch.cos(position * div_term)
+
+        self.emb = nn.Embedding(c_in, d_model)
+        self.emb.weight = nn.Parameter(w, requires_grad=False)
+
+    def forward(self, x):
+        return self.emb(x).detach()
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='fixed', freq='h'):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 4
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
+        if freq == 't':
+            self.minute_embed = Embed(minute_size, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+
+    def forward(self, x):
+        x = x.long()
+        minute_x = self.minute_embed(x[:, :, 4]) if hasattr(
+            self, 'minute_embed') else 0.
+        hour_x = self.hour_embed(x[:, :, 3])
+        weekday_x = self.weekday_embed(x[:, :, 2])
+        day_x = self.day_embed(x[:, :, 1])
+        month_x = self.month_embed(x[:, :, 0])
+
+        return hour_x + weekday_x + day_x + month_x + minute_x
+
+
+class TimeFeatureEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='timeF', freq='h'):
+        super(TimeFeatureEmbedding, self).__init__()
+
+        freq_map = {'h': 4, 't': 5, 's': 6,
+                    'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
+        d_inp = freq_map[freq]
+        self.embed = nn.Linear(d_inp, d_model, bias=False)
+
+    def forward(self, x):
+        return self.embed(x)
+
+
+class DataEmbedding(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        if x_mark is None:
+            x = self.value_embedding(x) + self.position_embedding(x)
+        else:
+            x = self.value_embedding(
+                x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
+        return self.dropout(x)
+
+
+class DataEmbedding_wo_pos(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding_wo_pos, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        if x_mark is None:
+            x = self.value_embedding(x)
+        else:
+            x = self.value_embedding(x) + self.temporal_embedding(x_mark)
+        return self.dropout(x)
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, d_model, patch_len, stride, dropout):
+        super(PatchEmbedding, self).__init__()
+        # Patching
+        self.patch_len = patch_len
+        self.stride = stride
+        self.padding_patch_layer = nn.ReplicationPad1d((0, stride))
+
+        # Backbone, Input encoding: projection of feature vectors onto a d-dim vector space
+        self.value_embedding = TokenEmbedding(patch_len, d_model)
+
+        # Positional embedding
+        self.position_embedding = PositionalEmbedding(d_model)
+
+        # Residual dropout
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # do patching
+        n_vars = x.shape[1]
+        x = self.padding_patch_layer(x)
+        x = x.unfold(dimension=-1, size=self.patch_len, step=self.stride)
+        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
+        # Input encoding
+        x = self.value_embedding(x) + self.position_embedding(x)
+        return self.dropout(x), n_vars
+
+class DataEmbedding_wo_time(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding_wo_time, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        x = self.value_embedding(x) + self.position_embedding(x)
+        return self.dropout(x)

code/TFC/main.py

@@ -1,42 +1,13 @@
-"""Updated implementation for TF-C -- Xiang Zhang, Jan 16, 2023"""
-
-import os
-import numpy as np
 from datetime import datetime
-import argparse
-from utils import _logger
+import sys 
+
 from model import *
-from dataloader import data_generator
+from dataloader import data_generator_motion
 from trainer import Trainer
 
-
-
-
-# Args selections
-start_time = datetime.now()
-parser = argparse.ArgumentParser()
-######################## Model parameters ########################
-home_dir = os.getcwd()
-parser.add_argument('--run_description', default='run1', type=str,
-                    help='Experiment Description')
-parser.add_argument('--seed', default=42, type=int, help='seed value')
-
-# 1. self_supervised pre_train; 2. finetune (itself contains finetune and test)
-parser.add_argument('--training_mode', default='fine_tune_test', type=str,
-                    help='pre_train, fine_tune_test')
-
-parser.add_argument('--pretrain_dataset', default='SleepEEG', type=str,
-                    help='Dataset of choice: SleepEEG, FD_A, HAR, ECG')
-parser.add_argument('--target_dataset', default='Epilepsy', type=str,
-                    help='Dataset of choice: Epilepsy, FD_B, Gesture, EMG')
-
-parser.add_argument('--logs_save_dir', default='../experiments_logs', type=str,
-                    help='saving directory')
-parser.add_argument('--device', default='cuda', type=str,
-                    help='cpu or cuda')
-parser.add_argument('--home_path', default=home_dir, type=str,
-                    help='Project home directory')
-args, unknown = parser.parse_known_args()
+sys.path.append("/home/ubuntu/projects/TFC-pretraining/code/")
+from config_files.MotionSample_Configs import Config as Configs
+from optimizers import get_optimizer
 
 with_gpu = torch.cuda.is_available()
 if with_gpu:
@@ -45,75 +16,57 @@
     device = torch.device("cpu")
 print('We are using %s now.' %device)
 
-pretrain_dataset = args.pretrain_dataset
-targetdata = args.target_dataset
-experiment_description = str(pretrain_dataset) + '_2_' + str(targetdata)
-
-
-method = 'TF-C'
-training_mode = args.training_mode
-run_description = args.run_description
-logs_save_dir = args.logs_save_dir
-os.makedirs(logs_save_dir, exist_ok=True)
-exec(f'from config_files.{pretrain_dataset}_Configs import Config as Configs')
+SAMPLE_RUN = True
+if SAMPLE_RUN: 
+  SAMPLE_COUNT, SAMPLE_COUNT_VAL = -1, -1
+  pretrain_dataset = 'MotionSample'
+else: 
+  SAMPLE_COUNT, SAMPLE_COUNT_VAL = 1000, 200
+  pretrain_dataset = 'Motion'
+
+training_mode = 'pre_train'
+MODEL_DIR = "/home/ubuntu/projects/TFC-pretraining/output/model"
+SAVE_MODEL_OR_CHECKPOINTS = 'checkpoints'
 configs = Configs()
 
-# # ##### fix random seeds for reproducibility ########
-SEED = args.seed
-torch.manual_seed(SEED)
-torch.backends.cudnn.deterministic = False
-torch.backends.cudnn.benchmark = False
-np.random.seed(SEED)
-#####################################################
-
-experiment_log_dir = os.path.join(logs_save_dir, experiment_description, run_description, training_mode + f"_seed_{SEED}_2layertransformer")
-# 'experiments_logs/Exp1/run1/train_linear_seed_0'
-os.makedirs(experiment_log_dir, exist_ok=True)
-
-# loop through domains
-counter = 0
-src_counter = 0
-
-
-# Logging
-log_file_name = os.path.join(experiment_log_dir, f"logs_{datetime.now().strftime('%d_%m_%Y_%H_%M_%S')}.log")
-# 'experiments_logs/Exp1/run1/train_linear_seed_0/logs_14_04_2022_15_13_12.log'
-logger = _logger(log_file_name)
-logger.debug("=" * 45)
-logger.debug(f'Pre-training Dataset: {pretrain_dataset}')
-logger.debug(f'Target (fine-tuning) Dataset: {targetdata}')
-logger.debug(f'Method:  {method}')
-logger.debug(f'Mode:    {training_mode}')
-logger.debug("=" * 45)
-
-# Load datasets
-sourcedata_path = f"../../datasets/{pretrain_dataset}"
-targetdata_path = f"../../datasets/{targetdata}"
-subset = True  # if subset= true, use a subset for debugging.
-train_dl, valid_dl, test_dl = data_generator(sourcedata_path, targetdata_path, configs, training_mode, subset = subset)
-logger.debug("Data loaded ...")
-
-# Load Model
-"""Here are two models, one basemodel, another is temporal contrastive model"""
-TFC_model = TFC(configs).to(device)
-classifier = target_classifier(configs).to(device)
-temporal_contr_model = None
-
-
-if training_mode == "fine_tune_test":
-    # load saved model of this experiment
-    load_from = os.path.join(os.path.join(logs_save_dir, experiment_description, run_description,
-    f"pre_train_seed_{SEED}_2layertransformer", "saved_models"))
-    print("The loading file path", load_from)
-    chkpoint = torch.load(os.path.join(load_from, "ckp_last.pt"), map_location=device)
-    pretrained_dict = chkpoint["model_state_dict"]
-    TFC_model.load_state_dict(pretrained_dict)
-
-model_optimizer = torch.optim.Adam(TFC_model.parameters(), lr=configs.lr, betas=(configs.beta1, configs.beta2), weight_decay=3e-4)
-classifier_optimizer = torch.optim.Adam(classifier.parameters(), lr=configs.lr, betas=(configs.beta1, configs.beta2), weight_decay=3e-4)
-
-# Trainer
-Trainer(TFC_model, model_optimizer, classifier, classifier_optimizer, train_dl, valid_dl, test_dl, device,
-        logger, configs, experiment_log_dir, training_mode)
-
-logger.debug(f"Training time is : {datetime.now()-start_time}")
+run_date = datetime.today().strftime('%Y-%m-%d')
+run_name = f"model_{run_date}"
+print(run_name)
+
+save_model_dir = MODEL_DIR
+
+# Load data 
+if SAMPLE_RUN: 
+  train_data_dir = '/home/ubuntu/projects/TFC-pretraining/data/train_data_sample.npy'
+  train_motion_names_dir = '/home/ubuntu/projects/TFC-pretraining/data/train_motion_names_sample.parquet'
+  train_fft_dir = '/home/ubuntu/projects/TFC-pretraining/data/train_fft_sample.npy'
+  val_data_dir = '/home/ubuntu/projects/TFC-pretraining/data/val_data_sample.npy'
+  val_motion_names_dir = '/home/ubuntu/projects/TFC-pretraining/data/val_motion_names_sample.parquet'
+  val_fft_dir = '/home/ubuntu/projects/TFC-pretraining/data/val_fft_sample.npy'
+else: 
+  train_data_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0512_TFC/fft_features/a6415950/train_data.npy'
+  train_motion_names_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0488_motion_periods/ts2vec_features/a6415950/train_motion_names.parquet'
+  train_fft_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0512_TFC/fft_features/a6415950/train_fft.npy'
+  val_data_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0512_TFC/fft_features/a6415950/val_data.npy'
+  val_motion_names_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0488_motion_periods/ts2vec_features/a6415950/val_motion_names.parquet'
+  val_fft_dir = '/dbfs/FileStore/fengdiguo/PDS184_driver_no_distraction/PDS0512_TFC/fft_features/a6415950/val_fft.npy'
+
+taindata_path_set = {'train_data_dir': train_data_dir, 
+                     'train_motion_names_dir': train_motion_names_dir, 
+                     'train_fft_dir': train_fft_dir}
+valdata_path_set = {'val_data_dir': val_data_dir, 
+                    'val_motion_names_dir': val_motion_names_dir, 
+                    'val_fft_dir': val_fft_dir}
+
+train_loader, val_loader, train_motion_names, val_motion_names = data_generator_motion(taindata_path_set, valdata_path_set, configs, training_mode, SAMPLE_COUNT, SAMPLE_COUNT_VAL)
+
+"""Load model"""
+# model = TFC(configs).to(device)
+# model_optimizer = torch.optim.Adam(model.parameters(), lr=configs.lr, betas=(configs.beta1, configs.beta2), weight_decay=3e-4)
+model = gpt4ts(configs).to(device)
+optim_class = get_optimizer(configs.optimizer)
+model_optimizer = optim_class(model.parameters(), lr=configs.gpt_lr, weight_decay=0)
+
+
+"""Model training"""
+train_loss, val_loss = Trainer(model, model_optimizer, train_loader, val_loader, device, configs,training_mode, save_model_dir, save_model_or_checkpoints=SAVE_MODEL_OR_CHECKPOINTS)