mirror of
https://github.com/Azure/MachineLearningNotebooks.git
synced 2025-12-20 01:27:06 -05:00
118 lines
4.9 KiB
Python
118 lines
4.9 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
# Licensed under the MIT license.
|
|
# Adapted from:
|
|
# https://github.com/Microsoft/CNTK/blob/master/Examples/Image/Classification/ConvNet/Python/ConvNet_MNIST.py
|
|
# ====================================================================
|
|
"""Train a CNN model on the MNIST dataset via distributed training."""
|
|
|
|
from __future__ import print_function
|
|
import numpy as np
|
|
import os
|
|
import cntk as C
|
|
import argparse
|
|
from cntk.train.training_session import CheckpointConfig, TestConfig
|
|
|
|
|
|
def create_reader(path, is_training, input_dim, label_dim, total_number_of_samples):
|
|
"""Define the reader for both training and evaluation action."""
|
|
return C.io.MinibatchSource(C.io.CTFDeserializer(path, C.io.StreamDefs(
|
|
features=C.io.StreamDef(field='features', shape=input_dim),
|
|
labels=C.io.StreamDef(field='labels', shape=label_dim)
|
|
)), randomize=is_training, max_samples=total_number_of_samples)
|
|
|
|
|
|
def convnet_mnist(max_epochs, output_dir, data_dir, debug_output=False, epoch_size=60000, minibatch_size=64):
|
|
"""Creates and trains a feedforward classification model for MNIST images."""
|
|
image_height = 28
|
|
image_width = 28
|
|
num_channels = 1
|
|
input_dim = image_height * image_width * num_channels
|
|
num_output_classes = 10
|
|
|
|
# Input variables denoting the features and label data
|
|
input_var = C.ops.input_variable((num_channels, image_height, image_width), np.float32)
|
|
label_var = C.ops.input_variable(num_output_classes, np.float32)
|
|
|
|
# Instantiate the feedforward classification model
|
|
scaled_input = C.ops.element_times(C.ops.constant(0.00390625), input_var)
|
|
|
|
with C.layers.default_options(activation=C.ops.relu, pad=False):
|
|
conv1 = C.layers.Convolution2D((5, 5), 32, pad=True)(scaled_input)
|
|
pool1 = C.layers.MaxPooling((3, 3), (2, 2))(conv1)
|
|
conv2 = C.layers.Convolution2D((3, 3), 48)(pool1)
|
|
pool2 = C.layers.MaxPooling((3, 3), (2, 2))(conv2)
|
|
conv3 = C.layers.Convolution2D((3, 3), 64)(pool2)
|
|
f4 = C.layers.Dense(96)(conv3)
|
|
drop4 = C.layers.Dropout(0.5)(f4)
|
|
z = C.layers.Dense(num_output_classes, activation=None)(drop4)
|
|
|
|
ce = C.losses.cross_entropy_with_softmax(z, label_var)
|
|
pe = C.metrics.classification_error(z, label_var)
|
|
|
|
# Load train data
|
|
reader_train = create_reader(os.path.join(data_dir, 'Train-28x28_cntk_text.txt'), True,
|
|
input_dim, num_output_classes, max_epochs * epoch_size)
|
|
# Load test data
|
|
reader_test = create_reader(os.path.join(data_dir, 'Test-28x28_cntk_text.txt'), False,
|
|
input_dim, num_output_classes, C.io.FULL_DATA_SWEEP)
|
|
|
|
# Set learning parameters
|
|
lr_per_sample = [0.001] * 10 + [0.0005] * 10 + [0.0001]
|
|
lr_schedule = C.learning_parameter_schedule_per_sample(lr_per_sample, epoch_size=epoch_size)
|
|
mms = [0] * 5 + [0.9990239141819757]
|
|
mm_schedule = C.learners.momentum_schedule_per_sample(mms, epoch_size=epoch_size)
|
|
|
|
# Instantiate the trainer object to drive the model training
|
|
local_learner = C.learners.momentum_sgd(z.parameters, lr_schedule, mm_schedule)
|
|
progress_printer = C.logging.ProgressPrinter(
|
|
tag='Training',
|
|
rank=C.train.distributed.Communicator.rank(),
|
|
num_epochs=max_epochs,
|
|
)
|
|
|
|
learner = C.train.distributed.data_parallel_distributed_learner(local_learner)
|
|
trainer = C.Trainer(z, (ce, pe), learner, progress_printer)
|
|
|
|
# define mapping from reader streams to network inputs
|
|
input_map_train = {
|
|
input_var: reader_train.streams.features,
|
|
label_var: reader_train.streams.labels
|
|
}
|
|
|
|
input_map_test = {
|
|
input_var: reader_test.streams.features,
|
|
label_var: reader_test.streams.labels
|
|
}
|
|
|
|
C.logging.log_number_of_parameters(z)
|
|
print()
|
|
|
|
C.train.training_session(
|
|
trainer=trainer,
|
|
mb_source=reader_train,
|
|
model_inputs_to_streams=input_map_train,
|
|
mb_size=minibatch_size,
|
|
progress_frequency=epoch_size,
|
|
checkpoint_config=CheckpointConfig(frequency=epoch_size,
|
|
filename=os.path.join(output_dir, "ConvNet_MNIST")),
|
|
test_config=TestConfig(reader_test, minibatch_size=minibatch_size,
|
|
model_inputs_to_streams=input_map_test)
|
|
).train()
|
|
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--num_epochs', help='Total number of epochs to train', type=int, default='40')
|
|
parser.add_argument('--output_dir', help='Output directory', required=False, default='outputs')
|
|
parser.add_argument('--data_dir', help='Directory with training data')
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
convnet_mnist(args.num_epochs, args.output_dir, args.data_dir)
|
|
|
|
# Must call MPI finalize when process exit without exceptions
|
|
C.train.distributed.Communicator.finalize()
|