-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
100 lines (82 loc) · 4.64 KB
/
main.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
import os
import sys
import logging
import pandas as pd
from functools import partial
from sklearn.model_selection import StratifiedKFold
import torch
from torch import optim, nn
from torch.utils.data import DataLoader
from models import Temp
from data import DataSet
from trainer import Trainer
from config import get_args
from lr_scheduler import get_sch
from utils import seed_everything, handle_unhandled_exception, save_to_json
if __name__ == "__main__":
args = get_args()
seed_everything(args.seed) #fix seed
device = torch.device('cuda:0') #use cuda:0
if args.continue_train > 0:
result_path = args.continue_from_folder
else:
result_path = os.path.join(args.result_path, args.model +'_'+str(len(os.listdir(args.result_path))))
os.makedirs(result_path)
logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger()
logger.addHandler(logging.FileHandler(os.path.join(result_path, 'log.log')))
logger.info(args)
save_to_json(vars(args), os.path.join(result_path, 'config.json'))
sys.excepthook = partial(handle_unhandled_exception,logger=logger)
train_data = pd.read_csv(args.train)
train_data['path'] = train_data['path'].apply(lambda x: os.path.join(args.path, x))
test_data = pd.read_csv(args.test)
test_data['path'] = test_data['path'].apply(lambda x: os.path.join(args.path, x))
#fix path based on the data dir
input_size = None
output_size = None
prediction = pd.read_csv(args.submission)
output_index = [f'{i}' for i in range(0, output_size)]
stackking_input = pd.DataFrame(columns = output_index, index=range(len(train_data))) #dataframe for saving OOF predictions
if args.continue_train > 0:
prediction = pd.read_csv(os.path.join(result_path, 'sum.csv'))
test_result = prediction[output_index].values
stackking_input = pd.read_csv(os.path.join(result_path, f'for_stacking_input.csv'))
skf = StratifiedKFold(n_splits=args.cv_k, random_state=args.seed, shuffle=True).split(train_data['path'], train_data['label']) #Using StratifiedKFold for cross-validation
for fold, (train_index, valid_index) in enumerate(skf): #by skf every fold will have similar label distribution
if args.continue_train > fold+1:
logger.info(f'skipping {fold+1}-fold')
continue
seed_everything(args.seed) #fix seed
fold_result_path = os.path.join(result_path, f'{fold+1}-fold')
os.makedirs(fold_result_path)
fold_logger = logger.getChild(f'{fold+1}-fold')
fold_logger.handlers.clear()
fold_logger.addHandler(logging.FileHandler(os.path.join(fold_result_path, 'log.log')))
fold_logger.info(f'start training of {fold+1}-fold')
#logger to log current n-fold output
kfold_train_data = train_data.iloc[train_index]
kfold_valid_data = train_data.iloc[valid_index]
train_dataset = DataSet(file_list=kfold_train_data['path'].values, label=kfold_train_data['label'].values)
valid_dataset = DataSet(file_list=kfold_valid_data['path'].values, label=kfold_valid_data['label'].values)
model = Temp(args).to(device) #make model based on the model name and args
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = get_sch(args.scheduler)(optimizer)
train_loader = DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, #pin_memory=True
)
valid_loader = DataLoader(
valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, #pin_memory=True
)
trainer = Trainer(
train_loader, valid_loader, model, loss_fn, optimizer, scheduler, device, args.patience, args.epochs, fold_result_path, fold_logger)
trainer.train() #start training
test_dataset = DataSet(file_list=test_data['path'].values, label=test_data['label'].values)
test_loader = DataLoader(
test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers
) #make test data loader
prediction[output_index] += trainer.test(test_loader) #softmax applied output; accumulate test prediction of current fold model
prediction.to_csv(os.path.join(result_path, 'sum.csv'), index=False)
stackking_input.loc[valid_index, output_index] = trainer.test(valid_loader) #use the validation data(hold out dataset) to make input for Stacking Ensemble model(out of fold prediction)
stackking_input.to_csv(os.path.join(result_path, f'for_stacking_input.csv'), index=False)