-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
92 lines (79 loc) · 1.9 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
import argparse
from torch import nn
from torch import optim
import checkpoint
import loader
import nnetwork
import trainer
def main():
ap = argparse.ArgumentParser(description='This script allow to train a pre-trained model')
ap.add_argument(
'data_dir',
default='./flowers'
)
ap.add_argument(
'--save_dir',
dest='save_dir',
default='./checkpoint.pth'
)
ap.add_argument(
'--learning_rate',
dest='learning_rate',
type=float,
default=0.001
)
ap.add_argument(
'--epochs',
dest='epochs',
type=int,
default=1
)
ap.add_argument(
'--arch',
dest='arch',
type=str,
default='vgg16'
)
ap.add_argument(
'--hidden_units',
dest='hidden_units',
type=int,
default=120
)
ap.add_argument(
'--gpu',
dest='gpu',
action='store_true'
)
args = ap.parse_args()
if args.data_dir:
# Load data
data = loader.load(args.data_dir)
# Setup Network
model = nnetwork.setup(
arch=args.arch,
hidden_units=args.hidden_units,
gpu=args.gpu
)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), args.learning_rate)
# Train Model
trainer.train(
model=model,
criterion=criterion,
optimizer=optimizer,
train_loader=data['train_loader'],
validation_loader=data['validation_loader'],
epochs=args.epochs,
gpu=args.gpu
)
# Save Checkpoint
checkpoint.save(
save_dir=args.save_dir,
model=model,
arch=args.arch,
hidden_units=args.hidden_units,
train_data=data['train_data']
)
if __name__ == "__main__":
main()