def train(dataset, origin_model, epochs=10):
"""
Train a model on the GTSRB dataset
"""
conv_model = build_model()
conv_model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
infos = origin_model.split('-')
poison_type = infos[1]
poison_loc = infos[2]
poison_size = int(infos[3])
filepath = "retrain_models/retrain"+"-{}".format(poison_type if poison_type else 'clean') \
+ '-%s-%d-{epoch:02d}-{val_acc:.2f}.hdf5'%(poison_loc,poison_size)
checkpoint = ModelCheckpoint(filepath,
monitor='val_acc',
verbose=1,
save_best_only=True,
mode='max')
callbacks_list = [checkpoint]
history = conv_model.fit(dataset.reprocessed_train_images,
dataset.reprocessed_train_labels,
callbacks=callbacks_list,
epochs=epochs,
validation_data=(dataset.test_images,
dataset.test_labels))
test_loss, test_acc = conv_model.evaluate(dataset.test_images,
dataset.test_labels,
verbose=2)
print("Test Loss: {}\nTest Acc: {}".format(test_loss, test_acc))
eval_badnet.evaluate_model(conv_model=conv_model)
def train(epochs=None,
poison_type=None,
poison_size=None,
poison_loc=None,
display=None):
"""
Train a model on the GTSRB dataset
"""
dataset = gtsrb_dataset.GTSRBDataset(poison_type=poison_type,
poison_size=poison_size,
poison_loc=poison_loc)
conv_model = build_model()
conv_model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
filepath = "output/badnet-{}".format(poison_type if poison_type else 'clean') \
+ '-{epoch:02d}-{val_acc:.2f}.hdf5'
checkpoint = ModelCheckpoint(filepath,
monitor='val_acc',
verbose=1,
save_best_only=True,
mode='max')
callbacks_list = [checkpoint]
history = conv_model.fit(dataset.train_images,
dataset.train_labels,
callbacks=callbacks_list,
epochs=epochs,
validation_data=(dataset.test_images,
dataset.test_labels))
if display:
plt.plot(history.history['acc'], label='accuracy')
plt.plot(history.history['val_acc'], label='val_accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim([0, 1])
plt.legend(loc='lower right')
plt.show()
test_loss, test_acc = conv_model.evaluate(dataset.test_images,
dataset.test_labels,
verbose=2)
print("Test Loss: {}\nTest Acc: {}".format(test_loss, test_acc))
eval_badnet.evaluate_model(conv_model=conv_model)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--epochs', default=1, type=int)
parser.add_argument('--checkpoint', type=str)
parser.add_argument('--poison-type', type=str)
parser.add_argument('--poison-loc', type=str)
parser.add_argument('--poison-size', type=int)
parser.add_argument('--train', action='store_true')
parser.add_argument('--eval', action='store_true')
parser.add_argument('--test-poison', action='store_true')
parser.add_argument('--display', action='store_true')
args = parser.parse_args()
if args.train:
train(epochs=args.epochs,
poison_type=args.poison_type,
poison_loc=args.poison_loc,
poison_size=args.poison_size,
display=args.display)
if args.eval:
eval_badnet.evaluate_model(checkpoint=args.checkpoint,
display=args.display,
conv_model=build_model())
if args.test_poison:
eval_badnet.test_poison(checkpoint=args.checkpoint,
conv_model=build_model(),
poison_type=args.poison_type,
poison_size=args.poison_size,
poison_loc=args.poison_loc)