-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathtrainer.py
193 lines (144 loc) · 7.11 KB
/
trainer.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
186
187
188
189
190
191
192
193
import torch
import torch.nn.functional as F
import datetime
import os
import collections
import numpy as np
import warnings
import sklearn.exceptions
warnings.filterwarnings("ignore", category=sklearn.exceptions.UndefinedMetricWarning)
warnings.simplefilter(action='ignore', category=FutureWarning)
from models import ecgTransForm
from dataloader import data_generator
from configs.data_configs import get_dataset_class
from configs.hparams import get_hparams_class
from utils import AverageMeter, to_device, _save_metrics, copy_Files
from utils import fix_randomness, starting_logs, save_checkpoint, _calc_metrics
class trainer(object):
def __init__(self, args):
# dataset parameters
self.dataset = args.dataset
self.seed_id = args.seed_id
self.device = torch.device(args.device)
# Exp Description
self.run_description = f"{args.run_description}_{datetime.datetime.now().strftime('%H_%M')}"
self.experiment_description = args.experiment_description
# paths
self.home_path = os.getcwd()
self.save_dir = os.path.join(os.getcwd(), "experiments_logs")
self.exp_log_dir = os.path.join(self.save_dir, self.experiment_description, self.run_description)
os.makedirs(self.exp_log_dir, exist_ok=True)
self.data_path = args.data_path
# Specify runs
self.num_runs = args.num_runs
# get dataset and base model configs
self.dataset_configs, self.hparams_class = self.get_configs()
# Specify hparams
self.hparams = self.hparams_class.train_params
def get_configs(self):
dataset_class = get_dataset_class(self.dataset)
hparams_class = get_hparams_class("supervised")
return dataset_class(), hparams_class()
def load_data(self, data_type):
self.train_dl, self.val_dl, self.test_dl, self.cw_dict = \
data_generator(self.data_path, data_type, self.hparams)
def calc_results_per_run(self):
acc, f1 = _calc_metrics(self.pred_labels, self.true_labels, self.dataset_configs.class_names)
return acc, f1
def train(self):
copy_Files(self.exp_log_dir) # save a copy of training files
self.metrics = {'accuracy': [], 'f1_score': []}
# fixing random seed
fix_randomness(int(self.seed_id))
# Logging
self.logger, self.scenario_log_dir = starting_logs(self.dataset, self.exp_log_dir, self.seed_id)
self.logger.debug(self.hparams)
# Load data
self.load_data(self.dataset)
model = ecgTransForm(configs=self.dataset_configs, hparams=self.hparams)
model.to(self.device)
# Average meters
loss_avg_meters = collections.defaultdict(lambda: AverageMeter())
self.optimizer = torch.optim.Adam(
model.parameters(),
lr=self.hparams["learning_rate"],
weight_decay=self.hparams["weight_decay"],
betas=(0.9, 0.99)
)
weights = [float(value) for value in self.cw_dict.values()]
# Now convert the list of floats to a numpy array, then to a PyTorch tensor
weights_array = np.array(weights).astype(np.float32) # Ensuring the correct dtype
weights_tensor = torch.tensor(weights_array).to(self.device)
self.cross_entropy = torch.nn.CrossEntropyLoss(weight=weights_tensor)
best_acc = 0
best_f1 = 0
# training..
for epoch in range(1, self.hparams["num_epochs"] + 1):
model.train()
for step, batches in enumerate(self.train_dl):
batches = to_device(batches, self.device)
data = batches['samples'].float()
labels = batches['labels'].long()
# ====== Source =====================
self.optimizer.zero_grad()
# Src original features
logits = model(data)
# Cross-Entropy loss
x_ent_loss = self.cross_entropy(logits, labels)
x_ent_loss.backward()
self.optimizer.step()
losses = {'Total_loss': x_ent_loss.item()}
for key, val in losses.items():
loss_avg_meters[key].update(val, self.hparams["batch_size"])
self.evaluate(model, self.val_dl)
tr_acc, tr_f1 = self.calc_results_per_run()
# logging
self.logger.debug(f'[Epoch : {epoch}/{self.hparams["num_epochs"]}]')
for key, val in loss_avg_meters.items():
self.logger.debug(f'{key}\t: {val.avg:2.4f}')
self.logger.debug(f'TRAIN: Acc:{tr_acc:2.4f} \t F1:{tr_f1:2.4f}')
# VALIDATION part
self.evaluate(model, self.val_dl)
ts_acc, ts_f1 = self.calc_results_per_run()
if ts_f1 > best_f1: # save best model based on best f1.
best_f1 = ts_f1
best_acc = ts_acc
save_checkpoint(self.exp_log_dir, model, self.dataset, self.dataset_configs, self.hparams, "best")
_save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "validation_best")
# logging
self.logger.debug(f'VAL : Acc:{ts_acc:2.4f} \t F1:{ts_f1:2.4f} (best: {best_f1:2.4f})')
self.logger.debug(f'-------------------------------------')
# LAST EPOCH
_save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "validation_last")
self.logger.debug("LAST EPOCH PERFORMANCE on validation set...")
self.logger.debug(f'Acc:{ts_acc:2.4f} \t F1:{ts_f1:2.4f}')
self.logger.debug(":::::::::::::")
# BEST EPOCH
self.logger.debug("BEST EPOCH PERFORMANCE on validation set ...")
self.logger.debug(f'Acc:{best_acc:2.4f} \t F1:{best_f1:2.4f}')
save_checkpoint(self.exp_log_dir, model, self.dataset, self.dataset_configs, self.hparams, "last")
# TESTING
print(" === Evaluating on TEST set ===")
self.evaluate(model, self.test_dl)
test_acc, test_f1 = self.calc_results_per_run()
_save_metrics(self.pred_labels, self.true_labels, self.exp_log_dir, "test_last")
self.logger.debug(f'Acc:{test_acc:2.4f} \t F1:{test_f1:2.4f}')
def evaluate(self, model, dataset):
model.to(self.device).eval()
total_loss_ = []
self.pred_labels = np.array([])
self.true_labels = np.array([])
with torch.no_grad():
for batches in dataset:
batches = to_device(batches, self.device)
data = batches['samples'].float()
labels = batches['labels'].long()
# forward pass
predictions = model(data)
# compute loss
loss = F.cross_entropy(predictions, labels)
total_loss_.append(loss.item())
pred = predictions.detach().argmax(dim=1) # get the index of the max log-probability
self.pred_labels = np.append(self.pred_labels, pred.cpu().numpy())
self.true_labels = np.append(self.true_labels, labels.data.cpu().numpy())
self.trg_loss = torch.tensor(total_loss_).mean() # average loss