-
Notifications
You must be signed in to change notification settings - Fork 2
/
train.py
executable file
·129 lines (106 loc) · 3.99 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
import yaml
import os
import argparse
import torch
import wandb
import numpy as np
import pandas as pd
import torch.nn as nn
import torch.optim as optim
from baseline import ModelBase, DataBase
from utils.run_train import run_train
def BinaryCrossEntropy(preds, labels):
probs = torch.sigmoid(preds)
labels = labels.float()
loss = torch.log(probs[:,0]+1e-10)*labels+torch.log((1-probs)[:,0]+1e-10)*(1-labels)
loss = -loss.mean()
return loss
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, required=True, help='config yaml path')
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--project', default='ClassTracing')
opt = parser.parse_args()
SEED = opt.seed
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
with open(opt.cfg, "r") as f:
cfg = yaml.safe_load(f)
print("Config training:")
for key, value in cfg.items():
print("{}: {}".format(key, value))
if not os.path.isdir(cfg['dir_checkpoint']):
print(f"Directory {cfg['dir_checkpoint']} does not exist")
os.makedirs(cfg['dir_checkpoint'])
print(f"Created {cfg['dir_checkpoint']}")
if not torch.cuda.is_available() and cfg['device'] == 'cuda':
print('Your device don\'t have cuda, please check or select cpu and retraining')
exit(0)
if cfg['device'] == 'cuda':
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# Build model
model_base = ModelBase(cfg['model'],
cfg['number_classes'],
device=cfg['device'])
model_base.build_model()
# Load data
df_train = pd.read_csv(cfg['df_train'])
df_val = pd.read_csv(cfg['df_val'])
data_base = DataBase(cfg['data'])
train_loader = data_base.get_dataloader(
df=df_train,
batch_size=cfg['batch_size'],
mode='train',
num_workers=cfg['num_worker']
)
val_loader = data_base.get_dataloader(
df=df_val,
batch_size=cfg['batch_size'],
mode='test',
num_workers=cfg['num_worker']
)
# Loss function and optimizer
if cfg['number_classes'] > 1:
loss_fn = nn.CrossEntropyLoss()
else:
loss_fn = BinaryCrossEntropy
optimizer = optim.AdamW(model_base.model.parameters(), lr=float(cfg['learning_rate']), betas=(0.9, 0.999))
wandb.init(project=opt.project, name=cfg['name_project'])
best_valid_acc = (-1.0) * float("Inf")
best_epoch = 0
for epoch in range(cfg['epoch']):
print(f"Epoch: {epoch+1}/{cfg['epoch']}")
train_metrics = run_train(
model=model_base.model,
dataloader=train_loader,
optimizer=optimizer,
criterion=loss_fn,
func_inference=model_base.inference,
mode='train'
)
val_metrics = run_train(
model=model_base.model,
dataloader=val_loader,
optimizer=optimizer,
criterion=loss_fn,
func_inference=model_base.inference,
mode='val'
)
print("Train metricc:", train_metrics)
print("Val metrics", val_metrics)
if cfg['save_each_epoch']:
torch.save(model_base.model.state_dict(),
cfg['dir_checkpoint'] + '/epoch_{}.pt'.format(epoch))
if best_valid_acc < val_metrics["val_acc"]:
best_valid_acc = val_metrics["val_acc"]
torch.save(model_base.model.state_dict(), cfg['dir_checkpoint'] + '/best.pt')
f = open(cfg['dir_checkpoint'] + '/best_epoch.txt', "w")
f.write(str(epoch))
print(f"Model saved to ==> {cfg['dir_checkpoint'] + '/best.pt'} at epoch {epoch}")
train_metrics.update(val_metrics)
wandb.log(train_metrics)
print(f'Finished training')