diff --git a/.gitignore b/.gitignore index 731bcb2..277c0b5 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,6 @@ datasets *.log *.pt *.zip +*.npy +*.pt +*.parquet diff --git a/code/TFC/embed.py b/code/TFC/embed.py new file mode 100644 index 0000000..5205d49 --- /dev/null +++ b/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) diff --git a/code/TFC/main.py b/code/TFC/main.py index fa5e523..88077bb 100644 --- a/code/TFC/main.py +++ b/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) diff --git a/code/TFC/main_old.py b/code/TFC/main_old.py new file mode 100644 index 0000000..fa5e523 --- /dev/null +++ b/code/TFC/main_old.py @@ -0,0 +1,119 @@ +"""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 +from model import * +from dataloader import data_generator +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() + +with_gpu = torch.cuda.is_available() +if with_gpu: + device = torch.device("cuda") +else: + 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') +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}") diff --git a/code/TFC/model.py b/code/TFC/model.py index e5de8c1..427dca9 100644 --- a/code/TFC/model.py +++ b/code/TFC/model.py @@ -1,6 +1,10 @@ from torch import nn import torch from torch.nn import TransformerEncoder, TransformerEncoderLayer +from transformers.models.gpt2.modeling_gpt2 import GPT2Model +from embed import DataEmbedding +import torch.nn.functional as F +from einops import rearrange """Two contrastive encoders""" class TFC(nn.Module): @@ -29,12 +33,14 @@ def __init__(self, configs): def forward(self, x_in_t, x_in_f): + # x_in_t: (batch_size, feat_dim, max_seq_len) + # x_in_f: (batch_size, feat_dim, max_seq_len) """Use Transformer""" - x = self.transformer_encoder_t(x_in_t) - h_time = x.reshape(x.shape[0], -1) + x = self.transformer_encoder_t(x_in_t) # (batch_size, feat_dim, max_seq_len) + h_time = x.reshape(x.shape[0], -1) # (batch_size, feat_dim * max_seq_len) """Cross-space projector""" - z_time = self.projector_t(h_time) + z_time = self.projector_t(h_time) # (batch_size, embedding) """Frequency-based contrastive encoder""" f = self.transformer_encoder_f(x_in_f) @@ -58,3 +64,78 @@ def forward(self, emb): emb = torch.sigmoid(self.logits(emb_flat)) pred = self.logits_simple(emb) return pred + + +"""GPT based transformers""" +class gpt4ts(nn.Module): + + def __init__(self, configs): + super(gpt4ts, self).__init__() + self.pred_len = 0 + self.seq_len = configs.max_seq_len + self.patch_size = configs.patch_size + self.stride = configs.stride + self.gpt_layers = configs.gpt_layers + self.feat_dim = configs.feat_dim + self.d_model = configs.d_model + + self.patch_num = (self.seq_len - self.patch_size) // self.stride + 1 + self.patch_num += 1 + + # time domain transformer + self.padding_patch_layer_t = nn.ReplicationPad1d((0, self.stride)) + self.enc_embedding_t = DataEmbedding(self.feat_dim * self.patch_size, configs.d_model, configs.dropout) + self.gpt2_t = GPT2Model.from_pretrained('gpt2', output_attentions=True, output_hidden_states=True) + # gpt2 has 12 layers + self.gpt2_t.h = self.gpt2_t.h[:self.gpt_layers] + for i, (name, param) in enumerate(self.gpt2_t.named_parameters()): + if 'ln' in name or 'wpe' in name: + param.requires_grad = True + else: + param.requires_grad = False + self.act = F.gelu + self.dropout_t = nn.Dropout(0.1) + self.ln_proj_t = nn.LayerNorm(configs.d_model * self.patch_num) + self.out_layer_t = nn.Linear(configs.d_model * self.patch_num, configs.embedding_len) + + # frequency domain transfomer + self.padding_patch_layer_f = nn.ReplicationPad1d((0, self.stride)) + self.enc_embedding_f = DataEmbedding(self.feat_dim * self.patch_size, configs.d_model, configs.dropout) + self.gpt2_f = GPT2Model.from_pretrained('gpt2', output_attentions=True, output_hidden_states=True) + self.gpt2_f.h = self.gpt2_f.h[:self.gpt_layers] + for i, (name, param) in enumerate(self.gpt2_f.named_parameters()): + if 'ln' in name or 'wpe' in name: + param.requires_grad = True + else: + param.requires_grad = False + self.act = F.gelu + self.dropout_f = nn.Dropout(0.1) + self.ln_proj_f = nn.LayerNorm(configs.d_model * self.patch_num) + self.out_layer_f = nn.Linear(configs.d_model * self.patch_num, configs.embedding_len) + + def forward(self, x_in_t, x_in_f): + # time domain + B_t, M_t, L_t = x_in_t.shape # B: batch_size, M: feat_dim, L: max_seq_len + input_t = self.padding_patch_layer_t(x_in_t) + input_t = input_t.unfold(dimension=-1, size=self.patch_size, step=self.stride) # (B, M, new_len, patch_size) + input_t = rearrange(input_t, 'b m n p -> b n (p m)') # (B, new_len, new_fea) + outputs_t = self.enc_embedding_t(input_t, None) # (B, new_len, d_model) + outputs_t = self.gpt2_t(inputs_embeds=outputs_t).last_hidden_state # (B, new_len, d_model) + h_time = outputs_t.reshape(outputs_t.shape[0], -1) # (B, new_len*d_model) + outputs_t = self.act(outputs_t).reshape(B_t, -1) # (B, new_len*d_model) + outputs_t = self.ln_proj_t(outputs_t) # (B, new_len*d_model) + z_time = self.out_layer_t(outputs_t) # (B, embedding_len) + + # freq domain + B_f, M_f, L_f = x_in_f.shape # B: batch_size, M: feat_dim, L: max_seq_len + input_f = self.padding_patch_layer_f(x_in_f) + input_f = input_f.unfold(dimension=-1, size=self.patch_size, step=self.stride) # (B, M, new_len, patch_size) + input_f = rearrange(input_f, 'b m n p -> b n (p m)') # (B, new_len, new_fea) + outputs_f = self.enc_embedding_f(input_f, None) # (B, new_len, d_model) + outputs_f = self.gpt2_f(inputs_embeds=outputs_f).last_hidden_state # (B, new_len, d_model) + h_freq = outputs_f.reshape(outputs_f.shape[0], -1) # (B, new_len*d_model) + outputs_f = self.act(outputs_f).reshape(B_f, -1) # (B, new_len*d_model) + outputs_f = self.ln_proj_f(outputs_f) # (B, new_len*d_model) + z_freq = self.out_layer_f(outputs_f) # (B, embedding_len) + + return h_time, z_time, h_freq, z_freq \ No newline at end of file diff --git a/code/TFC/optimizers.py b/code/TFC/optimizers.py new file mode 100644 index 0000000..c294a13 --- /dev/null +++ b/code/TFC/optimizers.py @@ -0,0 +1,263 @@ +import math +import torch +from torch.optim.optimizer import Optimizer + + +def get_optimizer(name): + + if name == "Adam": + return torch.optim.Adam + elif name == "RAdam": + return RAdam + + +# from https://github.com/LiyuanLucasLiu/RAdam/blob/master/radam/radam.py +class RAdam(Optimizer): + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, degenerated_to_sgd=True): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + + self.degenerated_to_sgd = degenerated_to_sgd + if isinstance(params, (list, tuple)) and len(params) > 0 and isinstance(params[0], dict): + for param in params: + if 'betas' in param and (param['betas'][0] != betas[0] or param['betas'][1] != betas[1]): + param['buffer'] = [[None, None, None] for _ in range(10)] + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, + buffer=[[None, None, None] for _ in range(10)]) + super(RAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(RAdam, self).__setstate__(state) + + def step(self, closure=None): + + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + raise RuntimeError('RAdam does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + # exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + # exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + + state['step'] += 1 + buffered = group['buffer'][int(state['step'] % 10)] + if state['step'] == buffered[0]: + N_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state['step'] + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + buffered[1] = N_sma + + # more conservative since it's an approximated value + if N_sma >= 5: + step_size = math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + elif self.degenerated_to_sgd: + step_size = 1.0 / (1 - beta1 ** state['step']) + else: + step_size = -1 + buffered[2] = step_size + + # more conservative since it's an approximated value + if N_sma >= 5: + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + denom = exp_avg_sq.sqrt().add_(group['eps']) + # p_data_fp32.addcdiv_(-step_size * group['lr'], exp_avg, denom) + p_data_fp32.addcdiv_(exp_avg, denom, value=-step_size * group['lr']) + p.data.copy_(p_data_fp32) + elif step_size > 0: + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + # p_data_fp32.add_(-step_size * group['lr'], exp_avg) + p_data_fp32.add_(exp_avg, alpha=-step_size * group['lr']) + p.data.copy_(p_data_fp32) + + return loss + + +class PlainRAdam(Optimizer): + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, degenerated_to_sgd=True): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + + self.degenerated_to_sgd = degenerated_to_sgd + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + + super(PlainRAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(PlainRAdam, self).__setstate__(state) + + def step(self, closure=None): + + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + raise RuntimeError('RAdam does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + state['step'] += 1 + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + + # more conservative since it's an approximated value + if N_sma >= 5: + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + step_size = group['lr'] * math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + denom = exp_avg_sq.sqrt().add_(group['eps']) + p_data_fp32.addcdiv_(-step_size, exp_avg, denom) + p.data.copy_(p_data_fp32) + elif self.degenerated_to_sgd: + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + step_size = group['lr'] / (1 - beta1 ** state['step']) + p_data_fp32.add_(-step_size, exp_avg) + p.data.copy_(p_data_fp32) + + return loss + + +class AdamW(Optimizer): + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, warmup=0): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + + defaults = dict(lr=lr, betas=betas, eps=eps, + weight_decay=weight_decay, warmup=warmup) + super(AdamW, self).__init__(params, defaults) + + def __setstate__(self, state): + super(AdamW, self).__setstate__(state) + + def step(self, closure=None): + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + denom = exp_avg_sq.sqrt().add_(group['eps']) + bias_correction1 = 1 - beta1 ** state['step'] + bias_correction2 = 1 - beta2 ** state['step'] + + if group['warmup'] > state['step']: + scheduled_lr = 1e-8 + state['step'] * group['lr'] / group['warmup'] + else: + scheduled_lr = group['lr'] + + step_size = scheduled_lr * math.sqrt(bias_correction2) / bias_correction1 + + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * scheduled_lr, p_data_fp32) + + p_data_fp32.addcdiv_(-step_size, exp_avg, denom) + + p.data.copy_(p_data_fp32) + + return loss diff --git a/code/config_files/MotionSample_Configs.py b/code/config_files/MotionSample_Configs.py index ef6d98c..4ba0b34 100644 --- a/code/config_files/MotionSample_Configs.py +++ b/code/config_files/MotionSample_Configs.py @@ -1,23 +1,30 @@ class Config(object): def __init__(self): - # model configs + # model configs for TFC self.input_channels = 8 # feature count self.transformer_nhead = 2 self.transformer_num_layers = 2 - self.embedding_len = 160 # final embedding len = embedding_len*2 - - # training configs - self.num_epoch = 2 - - # optimizer parameters + self.TSlength_aligned = 900 # sequence length 15Hz * 60second self.beta1 = 0.9 self.beta2 = 0.99 self.lr = 3e-4 # original lr: 3e-4 - # data parameters - self.drop_last = True + # model configs for gpt4ts + self.patch_size = 32 + self.stride = 16 + self.d_model = 768 # this is determined by gpt model. + self.dropout = 0.1 + self.gpt_layers = 2 + self.feat_dim = 8 # feature count + self.max_seq_len = 900 # seq lenth + self.optimizer = "RAdam" + self.gpt_lr = 0.001 + + # common configs + self.embedding_len = 160 # final embedding len = embedding_len*2 + self.num_epoch = 2 self.batch_size = 2 - self.TSlength_aligned = 900 # sequence length 15Hz * 60second + self.drop_last = True self.Context_Cont = Context_Cont_configs() self.TC = TC() @@ -26,8 +33,8 @@ def __init__(self): class augmentations(object): def __init__(self): - self.jitter_scale_ratio = 0.001 - self.jitter_ratio = 0.001 + self.jitter_scale_ratio = 1.1 + self.jitter_ratio = 0.8 self.max_seg = 5 diff --git a/code/config_files/Motion_Configs.py b/code/config_files/Motion_Configs.py index 6d5a75e..1dddc0d 100644 --- a/code/config_files/Motion_Configs.py +++ b/code/config_files/Motion_Configs.py @@ -1,23 +1,30 @@ class Config(object): def __init__(self): - # model configs + # model configs for TFC self.input_channels = 8 # feature count self.transformer_nhead = 2 self.transformer_num_layers = 2 - self.embedding_len = 160 # final embedding len = embedding_len*2 - - # training configs - self.num_epoch = 10 - - # optimizer parameters + self.TSlength_aligned = 900 # sequence length 15Hz * 60second self.beta1 = 0.9 self.beta2 = 0.99 - self.lr = 3e-4 # original lr: 3e-4 - - # data parameters + self.lr = 3e-4 # original lr: 3e-4 + + # model configs for gpt4ts + self.patch_size = 32 + self.stride = 16 + self.d_model = 768 # this is determined by gpt model. + self.dropout = 0.1 + self.gpt_layers = 6 + self.feat_dim = 8 # feature count + self.max_seq_len = 900 # seq lenth + self.optimizer = "RAdam" + self.gpt_lr = 0.001 + + # common configs + self.embedding_len = 160 # final embedding len = embedding_len*2 + self.num_epoch = 2 + self.batch_size = 2 self.drop_last = True - self.batch_size = 16 - self.TSlength_aligned = 900 # sequence length 15Hz * 60second self.Context_Cont = Context_Cont_configs() self.TC = TC() @@ -26,8 +33,8 @@ def __init__(self): class augmentations(object): def __init__(self): - self.jitter_scale_ratio = 0.001 - self.jitter_ratio = 0.001 + self.jitter_scale_ratio = 1.1 + self.jitter_ratio = 0.8 self.max_seg = 5