SecureTL/ecg_train_triplet_model.py at main · jtrpinto/SecureTL · GitHub

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
'''
Secure Triplet Loss Project Repository (https://github.com/jtrpinto/SecureTL)

File: ecg_train_triplet_model.py
- Used to train a model with ECG data, with the original triplet loss.

"Secure Triplet Loss: Achieving Cancelability and Non-Linkability in End-to-End Deep Biometrics"
João Ribeiro Pinto, Miguel V. Correia, and Jaime S. Cardoso
IEEE Transactions on Biometrics, Behavior, and Identity Science

joao.t.pinto@inesctec.pt  |  https://jtrpinto.github.io
'''

import os
import torch
import numpy as np
import pickle as pk
from models import TripletModel, TripletECGNetwork
from losses import TripletLoss
from dataset import TripletECGDataset
from trainer import train_triplet_model
from torch.utils.data import DataLoader


DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'

SAVE_MODEL = 'model_name'
TRAIN_DATA = 'ecg_train_data.pickle'

LEARN_RATE = 1e-4  # learning rate
REG = 0.001         # L2 regularization hyperparameter

N_EPOCHS = 250
BATCH_SIZE = 32
VALID_SPLIT = .2

print('Training model: ' + SAVE_MODEL)

# Preparing and dividing the dataset
trainset = TripletECGDataset(TRAIN_DATA)

dataset_size = len(trainset)  # number of samples in training + validation sets
indices = list(range(dataset_size))
split = int(np.floor(VALID_SPLIT * dataset_size))  # number of samples in validation set
np.random.seed(42)
np.random.shuffle(indices)
train_indices, valid_indices = indices[split:], indices[:split]

train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_indices)
valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(valid_indices)

train_loader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, sampler=train_sampler)
valid_loader = DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4, sampler=valid_sampler)

# Creating the network and the model
network = TripletECGNetwork().to(DEVICE)
model = TripletModel(network)
loss = TripletLoss(margin=1.0)
optimizer = torch.optim.Adam(model.parameters(), lr=LEARN_RATE, weight_decay=REG)

# Training the model
train_hist, valid_hist = train_triplet_model(model, loss, optimizer, train_loader, N_EPOCHS, BATCH_SIZE, DEVICE, patience=10, valid_loader=valid_loader, filename=SAVE_MODEL)