-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain_TSTCC_v1.py
185 lines (160 loc) · 5.69 KB
/
main_TSTCC_v1.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
import os
import shutil
import sys
from models.baselines.ml_baselines import get_baseline_performance
from sklearn.ensemble import AdaBoostClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from dataset.EMG_Gesture_v2 import EMGGestureDataModule
from dataset.Ninapro_DB5 import NinaproDB5DataModule
from models.TSTCC.lit_model import LitTSTCC
from configs.TSTCC_configs import Configs
import lightning.pytorch as pl
from lightning.pytorch.loggers import TensorBoardLogger
import torch
from lightning.pytorch import seed_everything
from shutil import copyfile
from preprocess.TSTCC_preprocess import TSTCCDataset
from utils.sampler import split_dataset
from utils.terminal_logger import TerminalLogger
from torch.utils.data import DataLoader
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--per_class_samples", type=int, default=100)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
# Initialization
config_dir = "configs/TSTCC_configs.py"
preprocess_dir = "preprocess/TSTCC_preprocess.py"
# config_dir = r"test_run/version_23/TFC_configs.py"
import_path = ".".join(config_dir.split(".")[0].split("/"))
print(f"from {import_path} import Configs")
exec(f"from {import_path} import Configs")
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
torch.set_float32_matmul_precision('medium')
configs = Configs()
configs.training_config.per_class_samples = args.per_class_samples
configs.training_config.seed = args.seed
configs.training_config.version = f"samples_{configs.training_config.per_class_samples}_pe_{configs.training_config.pretrain_epoch}_fe_{configs.training_config.finetune_epoch}_seed_{configs.training_config.seed}"
log_dir = os.path.join(
configs.training_config.log_save_dir,
configs.training_config.experiment_name,
configs.training_config.version,
"log"
)
if os.path.exists(log_dir):
shutil.rmtree(log_dir)
os.makedirs(log_dir)
seed_everything(configs.training_config.seed)
for fn in configs.training_config.bag_of_metrics.values():
fn.to(device)
# dataset = EMGGestureDataModule(
# dataset_type=TSTCCDataset,
# config=configs.dataset_config,
# )
dataset = NinaproDB5DataModule(
dataset_type=TSTCCDataset,
config=configs.dataset_config,
)
dataset.prepare_data()
pretrain_dataset = dataset.pretrain_dataset()
finetune_dataset = dataset.finetune_dataset()
finetune_train, val_and_test = split_dataset(
dataset=finetune_dataset,
num_samples=configs.training_config.per_class_samples,
shuffle=True,
)
finetune_val, finetune_test = split_dataset(
dataset=val_and_test,
num_samples=0.5,
shuffle=True,
)
lit_TSTCC = LitTSTCC(configs)
logger = TensorBoardLogger(
save_dir=configs.training_config.log_save_dir,
name=configs.training_config.experiment_name,
version=configs.training_config.version,
)
if "pretrain" in configs.training_config.mode:
lit_TSTCC.pretrain()
pretrain_loop = pl.Trainer(
deterministic=False,
max_epochs=configs.training_config.pretrain_epoch,
precision="16-mixed",
logger=logger,
log_every_n_steps=1,
)
pretrain_loop.fit(
model=lit_TSTCC,
train_dataloaders=DataLoader(
dataset=pretrain_dataset,
batch_size=configs.dataset_config.batch_size,
shuffle=True,
),
)
if "freeze" in configs.training_config.mode:
lit_TSTCC.freeze_encoder()
if "finetune" in configs.training_config.mode:
lit_TSTCC.finetune()
finetune_loop = pl.Trainer(
deterministic=False,
max_epochs=configs.training_config.finetune_epoch,
precision="16-mixed",
logger=logger,
log_every_n_steps=1,
enable_checkpointing=False,
)
finetune_loop.fit(
model=lit_TSTCC,
train_dataloaders=DataLoader(
dataset=finetune_train,
batch_size=configs.dataset_config.batch_size,
shuffle=True,
),
# val_dataloaders=DataLoader(
# dataset=finetune_val,
# batch_size=configs.dataset_config.batch_size,
# shuffle=True,
# pin_memory=True,
# )
)
finetune_loop.test(
model=lit_TSTCC,
dataloaders=DataLoader(
dataset=finetune_test,
batch_size=configs.dataset_config.batch_size,
shuffle=True,
),
)
# for fn in configs.training_config.bag_of_metrics.values():
# fn.to("cpu")
# baseline_model = [
# DecisionTreeClassifier(),
# KNeighborsClassifier(),
# AdaBoostClassifier(),
# ]
#
# get_baseline_performance(
# models=baseline_model,
# train_data=finetune_train,
# test_data=finetune_test,
# metrics=configs.training_config.bag_of_metrics,
# )
config_save_dir = os.path.join(logger.log_dir, config_dir.split("/")[1])
print(f"Saving config file at {config_save_dir}")
print(copyfile(config_dir, config_save_dir))
preprocess_save_dir = os.path.join(logger.log_dir, preprocess_dir.split("/")[1])
print(f"Saving preprocess file at {preprocess_save_dir}")
print(copyfile(preprocess_dir, preprocess_save_dir))
# input("Press any key to end the process")
# log_save_dir = "log.txt"
# print(f"Saving log file at {log_save_dir}")
# if os.path.exists(os.path.join(logger.log_dir, "log.txt")):
# os.remove(os.path.join(logger.log_dir, "log.txt"))
# print(copyfile(log_save_dir, os.path.join(logger.log_dir, "log.txt")))
# err_log_save_dir = "err_log.txt"
# print(f"Saving err log file at {err_log_save_dir}")
# if os.path.exists(os.path.join(logger.log_dir, "err_log.txt")):
# os.remove(os.path.join(logger.log_dir, "err_log.txt"))
# print(copyfile(log_save_dir, os.path.join(logger.log_dir, "err_log.txt")))