-
Notifications
You must be signed in to change notification settings - Fork 2
/
train_alexnet.py
66 lines (56 loc) · 2.27 KB
/
train_alexnet.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
# set the matplotlib backend so figures can be saved in the background
import matplotlib
matplotlib.use("Agg")
# import packages
from config import dogs_vs_cats_config as config
from pipeline.preprocessing import ImageToArrayPreprocessor
from pipeline.preprocessing import SimplePreprocessor
from pipeline.preprocessing import PatchPreprocessor
from pipeline.preprocessing import MeanPreprocessor
from pipeline.callbacks import TrainingMonitor
from pipeline.io import HDF5DatasetGenerator
from pipeline.nn.conv import AlexNet
from keras.preprocessing.image import ImageDataGenerator
from keras.optimizers import Adam
import json
import os
# construct the training image generator for data augmentation
aug = ImageDataGenerator(rotation_range = 20, zoom_range = 0.15,
width_shift_range = 0.2, height_shift_range = 0.2, shear_range = 0.15,
horizontal_flip = True, fill_mode = "nearest")
# load the RGB means for the training set
means = json.loads(open(config.DATASET_MEAN).read())
# initialize the image preprocessors
sp = SimplePreprocessor(227, 227)
pp = PatchPreprocessor(227, 227)
mp = MeanPreprocessor(means["R"], means["G"], means["B"])
iap = ImageToArrayPreprocessor()
# initialize the training and validation dataset generators
trainGen = HDF5DatasetGenerator(config.TRAIN_HDF5, 128, aug = aug,
preprocessors = [pp, mp, iap], classes = 2)
valGen = HDF5DatasetGenerator(config.VAL_HDF5, 128,
preprocessors = [sp, mp, iap], classes = 2)
# initialize the optimizer
print("[INFO] compiling model...")
opt = Adam(lr = 1e-3)
model = AlexNet.build(width = 227, height = 227, depth = 3, classes = 2, reg = 0.0002)
model.compile(loss = "binary_crossentropy", optimizer = opt, metrics = ["accuracy"])
# construct the set of callbacks
path = os.path.sep.join([config.OUTPUT_PATH, "{}.png".format(os.getpid())])
callbacks = [TrainingMonitor(path)]
# train the network
model.fit_generator(
trainGen.generator(),
steps_per_epoch = trainGen.numImages // 128,
validation_data = valGen.generator(),
validation_steps = valGen.numImages // 128,
epochs = 75,
max_queue_size = 10,
callbacks = callbacks, verbose = 1
)
# save the model to file
print("[INFO] serializing model...")
model.save(config.MODEL_PATH, overwrite = True)
# close the HDF5 dataset
trainGen.close()
valGen.close()