-
Notifications
You must be signed in to change notification settings - Fork 26
/
Copy pathtrain.py
154 lines (122 loc) · 4.6 KB
/
train.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
import pandas as pd
import numpy as np
import os
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from tqdm import tqdm
from torch.optim import AdamW
import time
import argparse as ap
from model_trainer import ModelTrainer
from classifier import RadBertClassifier
from dataset import CTDataset
from cosine_annealing_warmup import CosineAnnealingWarmupRestarts
from torch.optim.lr_scheduler import ReduceLROnPlateau
parser = ap.ArgumentParser()
parser.add_argument(
'--dataset',
help='dataset',
required=True,
default=None)
parser.add_argument(
'--augment',
help='augmentation',
required=False,
default=0,
type=int)
parser.add_argument(
'--scheduler',
help='scheduler, cawr or rlop',
required=False,
default=None)
args = parser.parse_args()
augment = True if args.augment == 1 else False
def get_unique_folder(base_folder):
counter = 1
new_folder = base_folder
while os.path.exists(new_folder):
new_folder = f"{base_folder}{counter}"
counter += 1
return new_folder
save_path = 'Results_'+os.path.basename(args.dataset)
save_path = get_unique_folder(save_path)
os.mkdir(save_path)
print('Results will be saved to ',save_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using ",device)
if device == 'cuda':
n_gpu = torch.cuda.device_count()
print("Number of GPU available:{} --> {} \n".format(n_gpu,torch.cuda.get_device_name()))
df = pd.read_csv(os.path.join(args.dataset,'train.csv'))
df2 = pd.read_csv(os.path.join(args.dataset,'val.csv'))
print('average sentence length: ', df['report_text'].str.split().str.len().mean())
print('stdev sentence length: ', df['report_text'].str.split().str.len().std())
cols = df.columns
label_cols = list(cols[2:])
num_labels = len(label_cols)
print('Label columns: ', label_cols)
# Create dataloader
dataloaders = {}
max_length = 512
num_workers = 4
batch_size = 32
train_data = CTDataset(df, num_labels, label_cols, max_length, augment = args.augment)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size,
num_workers=num_workers, pin_memory=True)
dataloaders['train'] = train_dataloader
val_data = CTDataset(df2, num_labels, label_cols, max_length)
val_sampler = SequentialSampler(val_data)
val_dataloader = DataLoader(val_data, sampler=val_sampler, batch_size=batch_size,
num_workers=num_workers, pin_memory=True)
dataloaders['val'] = val_dataloader
model = RadBertClassifier(n_classes=num_labels)
model = model.to(device)
# setting custom optimization parameters
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'gamma', 'beta']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay_rate': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay_rate': 0.0}
]
lr_rate = 2e-5
optimizer = AdamW(optimizer_grouped_parameters,lr=lr_rate)
# optimizer = AdamW(model.parameters(),lr=2e-5) # Default optimization
epochs = 1000
w_steps = 50
cycle_step = 200
if args.scheduler == 'cawr':
scheduler = CosineAnnealingWarmupRestarts(optimizer,
first_cycle_steps=cycle_step,
cycle_mult = 1,
max_lr = lr_rate,
min_lr = 2e-7,
warmup_steps = w_steps)
elif args.scheduler =='rlop':
scheduler = ReduceLROnPlateau(optimizer, mode='min', min_lr = 2e-6, factor=0.1, patience=25)
else:
scheduler = None
trainer = ModelTrainer(model,
dataloaders,
num_labels,
epochs,
optimizer,
scheduler,
device,
save_path,
label_cols)
start = time.time()
print('----------------------Starting Training----------------------')
print("Number of epochs: ",epochs)
cm, clf_report = trainer.launch_training()
finish = time.time()
print('---------------------------------------------------------------')
print('Training_complete')
print('Training time: ',finish-start)
print(clf_report)
with open(os.path.join(save_path, 'test_classification_report.txt'), 'w') as file:
file.write(clf_report)
# Save confusion matrix
np.save(os.path.join(save_path, 'test_confusion_matrix.npy'), cm)