Amazon Sagemaker モデルのパフォーマンスのリアルタイム監視に Amazon CloudWatch カスタムメトリックスを使用する

import sagemaker
from sagemaker.mxnet import MXNet
from mxnet import gluon
from sagemaker import get_execution_role
# Creates a SageMaker session to upload the training data
sagemaker_session = sagemaker.Session()

# Returns the IAM role that SageMaker uses to access Amazon S3 and other services.
role = get_execution_role()

# Download the training and testing data
gluon.data.vision.MNIST('./data/train', train=True)
gluon.data.vision.MNIST('./data/test', train=False)

# Upload the downloaded data to an S3 location 
input_data = sagemaker_session.upload_data(path='data', key_prefix='data/mnist')

# Creating an MXNet Estimator that wraps the training environment parameters. P2.xlarge is a GPU type instance and will use GPU verion of MXNet
mxnet_estimator = MXNet("training_function.py", role=role, train_instance_count=1, train_instance_type="ml.p2.xlarge",
          hyperparameters={'batch_size': 100, 'epochs': 20, 'learning_rate': 0.1,'momentum': 0.9,'log_interval': 100})

# start the training process.
mxnet_estimator.fit(input_data)

def train(channel_input_dirs, model_dir, hyperparameters, hosts, num_gpus, **kwargs):
    # SageMaker passes num_cpus, num_gpus and other args we can use to tailor training to the current container environment.
    # The training will use GPU if the instance type is a GPU instance otherwise, it will use CPU.
    ctx = mx.gpu() if num_gpus > 0 else mx.cpu()

    num_examples = 70000
    best_accuracy = 0.0

    # retrieve the hyperparameters we set in notebook (with some defaults)
    batch_size = hyperparameters.get('batch_size', 100)
    epochs = hyperparameters.get('epochs', 20)
    learning_rate = hyperparameters.get('learning_rate', 0.1)
    momentum = hyperparameters.get('momentum', 0.9)
    log_interval = hyperparameters.get('log_interval', 100)

    # load training and validation data
    # we use the gluon.data.vision.MNIST class because of its built-in MNIST pre-processing logic,
    # but point it at the location where SageMaker placed the data files, so it doesn't download them again.
    training_dir = channel_input_dirs['training']
    train_data = get_train_data(training_dir + '/train', batch_size)
    val_data = get_val_data(training_dir + '/test', batch_size)

    # define the network
    net = nn.Sequential()
    with net.name_scope():
      net.add(nn.Dense(128, activation='relu'))
      net.add(nn.Dense(64, activation='relu'))
      net.add(nn.Dense(10))

    # Collect all parameters from net and its children, then initialize them.
    net.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)

    if len(hosts) == 1:
        kvstore = 'device' if num_gpus > 0 else 'local'
    else:
        kvstore = 'dist_device_sync' if num_gpus > 0 else 'dist_sync'

    trainer = gluon.Trainer(net.collect_params(), 'sgd',
                            {'learning_rate': learning_rate, 'momentum': momentum},
                            kvstore=kvstore)

    # Instantiate the Accuracy and Cross Entropy Loss metric classes.
    # The later Computes the softmax cross entropy loss.
    metric = mx.metric.Accuracy()
    loss = gluon.loss.SoftmaxCrossEntropyLoss()

    # Instatiate the CloudWatch Class for Evaluation Metrics.
    CWMetrics = CWEvalMetrics(region=region, model_name=model_name)   
    
    for epoch in range(epochs):
        # reset data iterator and metric at begining of epoch.
        metric.reset()
        cumulative_loss = 0
        btic = time.time()

        for i, (data, label) in enumerate(train_data):
            # Copy data to ctx if necessary
            data = data.as_in_context(ctx)
            label = label.as_in_context(ctx)

            # Start recording computation graph with record() section.
            # Recorded graphs can then be differentiated with backward.
            with autograd.record():
                output = net(data)
                L = loss(output, label)
                L.backward()
            # Get the sum of the Loss for each Iteration
            cumulative_loss += nd.sum(L).asscalar()
            trainer.step(data.shape[0])

            # update metric at last.
            metric.update([label], [output])
            
            if i % log_interval == 0 and i > 0:
                name, acc = metric.get()
                print('Training Loss: %f' % float(cumulative_loss / num_examples))
                print('[Epoch %d Batch %d] Training: %s=%f, %f samples/s' %
                      (epoch, i, name, acc, batch_size / (time.time() - btic)))

                # Send the collected training accuracy and loss to CloudWatch
                CWMetrics.CW_eval(model_name, is_training=True, Accuracy=acc * 100, Loss=float(cumulative_loss / num_examples * 100), hyperparameters=hyperparameters)
                              
            btic = time.time()

        name, acc = metric.get()
        print('[Epoch %d] Training: %s=%f' % (epoch, name, acc))

        name, val_acc = test(ctx, net, val_data)
        print('[Epoch %d] Validation: %s=%f' % (epoch, name, val_acc))

        #Saving the model if this is the best validation accuracy in the train loop
        if val_acc > best_accuracy:
    			symbol = net(mx.sym.var('data'))
    			symbol.save('%s/model.json' % model_dir)
    			net.collect_params().save('{}/model-{:0>4}.params'.format(model_dir, epoch))
    			best_accuracy = val_ac

        # Send the collected training accuracy and loss to CloudWatch
        CWMetrics.CW_eval(model_name, is_training=False, Accuracy=val_acc * 100, hyperparameters=hyperparameters)
         
        # Create Cloudwatch Dashboard with the generated metrics.
       		if epoch == 1:
	   CWMetrics.create_dashboard('MNIST_Dashboard')
    return net

class CWEvalMetrics():
    # initialize the region and the model name with the class instantiation
    def __init__(self, region='us-east-1', model_name='CW-example'):
        self.region = region
        self.model_name = model_name

    # A function to send the training evaluation metrics
    # the metric_type parameters will determine whether the data sent is for training or validation.

    def CW_eval(self, model_name, is_training,  **kwargs):
        # collecting the loss and accuracy values
        loss = kwargs.get('Loss', 0)
        accuracy = kwargs.get('Accuracy')

        # determine if the passed values are for training or validation
        if is_training:
            metric_type = 'Training'
        else:
            metric_type = 'Validation'

        # Collecting the hyperparameters to be used as the metrics dimensions
        hyperparameter = kwargs.get('hyperparameters')
        optimizer = str(hyperparameter.get('optimizer'))
        epochs = str(hyperparameter.get('epochs'))
        learning_rate = str(hyperparameter.get('learning_rate'))
        response = client.put_metric_data(
            Namespace='/aws/sagemaker/' + model_name,
            MetricData=[
                {
                    'MetricName': metric_type + ' Accuracy',
                    'Dimensions': [
                  { 'Name': 'Model Name', 'Value': model_name },
                  { 'Name': 'Learning Rate', 'Value': learning_rate },
                  { 'Name': 'Optimizer', 'Value': optimizer },
                  { 'Name': 'Epochs', 'Value': epochs},
                    ],
                    'Value': accuracy,
                    'Unit': "Percent",
                    'StorageResolution': 1
                },
                {
                    'MetricName': metric_type + ' Loss',
                    'Dimensions': [
                  { 'Name': 'Model Name', 'Value': model_name },
                  { 'Name': 'Learning Rate', 'Value': learning_rate },
                  { 'Name': 'Optimizer', 'Value': optimizer },
                  { 'Name': 'Epochs', 'Value': epochs},
                    ],
                    'Value': loss,
                    'Unit': "Percent",
                    'StorageResolution': 1
                },
            ]
        )
        return response

# A function to create a dashboard with the above training metrics.
def create_dashboard(self, db_name, **kwargs):
    hyperparameter = kwargs.get('hyperparameters')
    job_name = str(hyperparameter.get('sagemaker_job_name'))
    optimizer = str(hyperparameter.get('optimizer'))
    epochs = str(hyperparameter.get('epochs'))
    lr = str(hyperparameter.get('learning_rate'))

    # The dashboard body has the property of the dashboard in JSON format
    dashboard_body = '{"widgets":[{"type":"metric","x":0,"y":3,"width":18,"height":9,"properties":{"view":"timeSeries","stacked":false,"metrics":[["/aws/sagemaker/' + self.model_name + '","Training Loss","Model Name","' + self.model_name + '","Epochs","' + epochs + '","Optimizer","' + optimizer + '","Learning Rate","' + lr + '"],[".","Training Accuracy",".",".",".",".",".",".",".","."],[".","Validation Accuracy",".",".",".",".",".",".",".","."]],"region":"' + self.region + '","period":30}},{"type":"metric","x":0,"y":0,"width":18,"height":3,"properties":{"view":"singleValue","metrics":[["/aws/sagemaker/' + self.model_name + '","Training Loss","Model Name","' + self.model_name + '","Epochs","' + epochs + '","Optimizer","' + optimizer + '","Learning Rate","' + lr + '"],[".","Training Accuracy",".",".",".",".",".",".",".","."],[".","Validation Accuracy",".",".",".",".",".",".",".","."]],"region":"' + self.region + '","period":30}}]}'

    response = client.put_dashboard(DashboardName=db_name, DashboardBody=dashboard_body)
    return response

# when creating an alarm that includes the 'Job Name' as one of the dimensions, then the CloudWatch
# metric created also needs to have 'Job Name' as one of the dimensions
def create_alarm(self, model_name, **kwargs):
    hyperparameter = kwargs.get('hyperparameters')
    job_name = str(hyperparameter.get('sagemaker_job_name'))
    optimizer = str(hyperparameter.get('optimizer'))
    epochs = str(hyperparameter.get('epochs'))
    learning_rate = str(hyperparameter.get('learning_rate'))
    namespace = '/aws/sagemaker/' + model_name
    client.put_metric_alarm(
        AlarmName='Overfitting Alarm',
        ComparisonOperator='LessThanOrEqualToThreshold',
        EvaluationPeriods=10,
        DatapointsToAlarm=10,
        MetricName='Validation Accuracy',
        Namespace=namespace,
        Period=30,
        Statistic='Maximum',
        Threshold=90.0,
        ActionsEnabled=True,
        AlarmActions=[
            'arn:aws:sns:us-west-2:<account-id>:Sagemaker-Notification-Emails',
        ],
        AlarmDescription='Alarm when the model is Overfitting',
        Dimensions=[
            { 'Name': 'Model Name', 'Value': model_name },
            { 'Name': 'Job Name', 'Value': job_name },
            { 'Name': 'Learning Rate', 'Value': learning_rate },
            { 'Name': 'Optimizer', 'Value': optimizer },
            { 'Name': 'Epochs', 'Value': epochs},
        ],
        Unit='Percent'
    )

import json
import boto3


def lambda_handler(event, context):
    # Get the job name from the event
    Trigger = event["Records"][0]["Sns"]["Message"]
    Job = json.loads(Trigger)

    # The index of the job name element will be the same as the one listed in the create_alarm() function
    Job = Job["Trigger"]["Dimensions"][1]["value"]

    # stop the training job if overfitting is detected
    client = boto3.client('sagemaker')
    response = client.stop_training_job(TrainingJobName=Job)

    return response