-
Notifications
You must be signed in to change notification settings - Fork 3
/
HCNN.py
87 lines (71 loc) · 4.46 KB
/
HCNN.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
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
import pytorch_lightning as pl
from homological_models import *
class HCNN:
def __init__(self, X_train, X_val, X_test, y_train, y_val, y_test,
n_filters_l1, n_filters_l2, tmfg_repetitions, tmfg_confidence,
tmfg_similarity, dropout_rate, max_epochs=params.MAX_EPOCHS, T=1, root_folder=None, filtering_type=None, seed=None):
self.number_of_selected_features, self.shape_tetrahedra, self.shape_triangles, self.shape_simplex, self.X_train, self.X_val, self.X_test, self.y_train, self.y_val, self.y_test = h_input_transform(
X_train=X_train,
X_val=X_val,
X_test=X_test,
y_train=y_train,
y_val=y_val,
y_test=y_test,
tmfg_repetitions=tmfg_repetitions,
tmfg_confidence=tmfg_confidence,
tmfg_similarity=tmfg_similarity,
filtering_type=filtering_type)
self.T = T
self.n_filters_l1 = n_filters_l1
self.n_filters_l2 = n_filters_l2
self.last_layer_neurons = len(pd.Series(y_train).unique())
self.best_model = None
self.model = HCNN_model1D(T=self.T,
FILTERS_L1=self.n_filters_l1,
FILTERS_L2=self.n_filters_l2,
last_layer_neurons=self.last_layer_neurons,
NF_4=self.shape_tetrahedra,
NF_3=self.shape_triangles,
NF_2=self.shape_simplex,
dropout_rate=dropout_rate,)
# Set the global seed for reproducibility.
pl.seed_everything(seed)
# Set learning rate based on OpenAI's implementation.
self.model.set_lr(get_openai_lr(self.model))
self.early_stopping = EarlyStopping('val_loss', mode='min', patience=params.EARLY_STOPPING_PATIENCE)
self.checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor="val_loss", mode="min")
self.logger_folder = root_folder + 'logs'
if params.DEVICE == 'cuda':
self.trainer = Trainer(max_epochs=max_epochs,
accelerator='gpu',
enable_progress_bar=False,
enable_model_summary=False,
callbacks=[self.early_stopping, self.checkpoint_callback],
default_root_dir=self.logger_folder,
)
else:
self.trainer = Trainer(max_epochs=max_epochs,
enable_progress_bar=False,
enable_model_summary=False,
callbacks=[self.early_stopping, self.checkpoint_callback],
default_root_dir=self.logger_folder,
)
self.train_dataloader, self.val_dataloader, self.test_dataloader = prepare_dataloaders(self.X_train,
self.X_val,
self.X_test,
self.y_train,
self.y_val,
self.y_test,
batch_size=params.BATCH_SIZE)
def fit(self):
self.trainer.fit(self.model, train_dataloaders=self.train_dataloader, val_dataloaders=self.val_dataloader)
def evaluate(self):
self.trainer.validate(self.model, dataloaders=self.val_dataloader, verbose=False, ckpt_path='best')
return self.model.val_targets, self.model.val_preds, self.model.val_probs
def predict(self):
self.trainer.test(self.model, dataloaders=self.test_dataloader, verbose=False, ckpt_path='best')
return self.model.test_targets, self.model.test_preds, self.model.test_probs
def data_preparation_pipeline(self):
return self.X_train, self.X_val, self.X_test, self.y_train, self.y_val, self.y_test, self.model, self.number_of_selected_features