def make_tensor_model(staging_area_callback, num_classes):
model = ResNet50(input_tensor=staging_area_callback.input_tensor,
classes=num_classes, weights=None)
model.compile(optimizer='sgd', loss='categorical_crossentropy',
target_tensors=[staging_area_callback.target_tensor],
feed_dict=staging_area_callback.feed_dict,
fetches=staging_area_callback.extra_ops)
return model
num_classes = 1000
dataset_size = 1024
batch_size = 32
epochs = 5
x_train, y_train = create_synth_imagenet(224, dataset_size)
x_train = x_train.astype('float32')
y_train = y_train.astype('float32')
# last batch might be smaller
steps_per_epoch = int(math.ceil(len(x_train) / batch_size))
gauge = SamplesPerSec(batch_size)
staging_area_callback = StagingAreaCallback(x_train, y_train, batch_size)
print('training plain model:')
plain_model = make_plain_model(num_classes)
plain_model.fit(x_train, y_train, batch_size, epochs=epochs, callbacks=[gauge])
print('training pipelined model:')
pipelined_model = make_tensor_model(staging_area_callback, num_classes)
parser.add_argument('-b',
'--batch-size',
default=32,
type=int,
help='Batch size')
return parser.parse_args()
if __name__ == '__main__':
args = parse_args()
print('architecture:', args.arch)
print('number of GPUs:', args.gpus)
print('parameter server:', args.parameter_server)
print('parallelization method:', args.method)
print('epochs:', args.epochs)
print('batch_size:', args.batch_size)
dataset_size = 1024
if args.arch == 'inception3':
create_model = InceptionV3
X, y = create_synth_imagenet(299, dataset_size)
elif args.arch == 'resnet50':
create_model = ResNet50
X, y = create_synth_imagenet(224, dataset_size)
train(create_model, X, y, args.batch_size, args.epochs, args.gpus,
args.parameter_server, args.method)