-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
114 lines (94 loc) · 3.94 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
# Imports
import argparse
import data_helper
import model_helper
import numpy as np
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models
import warnings
warnings.filterwarnings("ignore")
# ----------------------------
parser = argparse.ArgumentParser(description='Build, train and save image classifier model')
parser.add_argument('data_dir',
#'--data_dir',
#action = 'store',
#dest = 'data_dir',
nargs='?',
type = str,
default = 'flowers',
help='directory of the dataset')
parser.add_argument('--gpu',
action='store_true',
dest = 'gpu',
default = False,
help = 'enable GPU mode')
parser.add_argument('--epochs',
action='store',
dest = 'epochs',
type = int,
default = 1,
help = 'number of epochs')
parser.add_argument('--arch',
action = 'store',
dest = 'arch',
type = str,
choices=["densenet121", "vgg16"],
default = 'densenet121',
help = 'pre-trained model architecture: densenet121 or vgg16')
parser.add_argument('--learning_rate',
action = 'store',
dest = 'learning_rate',
type = float,
default = 0.001,
help = 'learning rate')
args = parser.parse_args()
print("\n")
print("--------- Entered Arguments ----------")
print("Data directory : {!r}".format(args.data_dir))
print("Enable gpu : {!r}".format(args.gpu))
print("Epoch(s) : {!r}".format(args.epochs))
print("Model architecture : {!r}".format(args.arch))
print("Learning rate : {!r}".format(args.learning_rate))
print("--------------------------------------\n")
if __name__ == "__main__":
device = torch.device("cuda:0" if args.gpu else "cpu")
trainloader, validloader, testloader, num_categories, cat_to_name, class_to_idx = data_helper.load_imgs(args.data_dir)
model = model_helper.build_model(args, num_categories)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=args.learning_rate)
model = model_helper.train_model(model,
args.arch,
trainloader, validloader,
args.epochs,
args.learning_rate,
criterion,
optimizer,
device)
model.eval()
print("\nTesting ... ")
# Turn off gradients for validation
with torch.no_grad():
test_loss, test_accuracy = model_helper.evaluate_model(model, testloader, criterion, device)
print("Testing Loss: {:.3f}.. ".format(test_loss/len(testloader)),
"Testing Accuracy: {0:.2%} \n".format(test_accuracy/len(testloader)))
# Save model
print("\nSaving checkpoint ... \n")
checkpoint = {'arch': args.arch,
'input_size': model.num_features,
'hidden': model.hidden,
'output_size': num_categories,
'batch_size': model.batch_size,
'state_dict': model.state_dict(),
'optimizer_dict':optimizer.state_dict(),
'class_to_idx': class_to_idx,
'epoch': model.epochs,
'dropout_p': model.dropout_p,
'learning_rate': model.learning_rate,
'cat_to_name': cat_to_name,
}
model_helper.print_model_details(checkpoint)
torch.save(checkpoint, 'my_checkpoint.pth')
print("Model saved successfully.\n")