Package: sagemaker-training
SageMaker Training ToolkitTrain machine learning models within a Docker container using Amazon SageMaker.:books: BackgroundAmazon SageMakeris a fully managed service for data science and machine learning (ML) workflows.
You can use Amazon SageMaker to simplify the process of building, training, and deploying ML models.To train a model, you can include your training script and dependencies in aDocker containerthat runs your training code.
A container provides an effectively isolated environment, ensuring a consistent runtime and reliable training process.TheSageMaker Training Toolkitcan be easily added to any Docker container, making it compatible with SageMaker fortraining models.
If you use aprebuilt SageMaker Docker image for training, this library may already be included.For more information, see the Amazon SageMaker Developer Guide sections onusing Docker containers for training.:hammer_and_wrench: InstallationTo install this library in your Docker image, add the following line to yourDockerfile:RUNpip3installsagemaker-training:computer: UsageThe following are brief how-to guides.
For complete, working examples of custom training containers built with the SageMaker Training Toolkit, please seethe example notebooks.Create a Docker image and train a modelWrite a training script (eg.train.py).Define a container with a Dockerfilethat includes the training script and any dependencies.The training script must be located in the/opt/ml/codedirectory.
The environment variableSAGEMAKER_PROGRAMdefines which file inside the/opt/ml/codedirectory to use as the training entry point.
When training starts, the interpreter executes the entry point defined bySAGEMAKER_PROGRAM.
Python and shell scripts are both supported.FROMyourbaseimage:tag# install the SageMaker Training ToolkitRUNpip3installsagemaker-training# copy the training script inside the containerCOPYtrain.py/opt/ml/code/train.py# define train.py as the script entry pointENVSAGEMAKER_PROGRAMtrain.pyBuild and tag the Docker image.dockerbuild-tcustom-training-container.Use the Docker image to start a training job using theSageMaker Python SDK.fromsagemaker.estimatorimportEstimatorestimator=Estimator(image_name="custom-training-container",role="SageMakerRole",train_instance_count=1,train_instance_type="local")estimator.fit()To train a model using the image on SageMaker,push the image to ECRand start a SageMaker training job with the image URI.Pass arguments to the entry point using hyperparametersAny hyperparameters provided by the training job are passed to the entry point as script arguments.
The SageMaker Python SDK uses this feature to pass special hyperparameters to the training job, includingsagemaker_programandsagemaker_submit_directory.
The complete list of SageMaker hyperparameters is availablehere.Implement an argument parser in the entry point script. For example, in a Python script:importargparseif__name__=="__main__":parser=argparse.ArgumentParser()parser.add_argument("--learning-rate",type=int,default=1)parser.add_argument("--batch-size",type=int,default=64)parser.add_argument("--communicator",type=str)parser.add_argument("--frequency",type=int,default=20)args=parser.parse_args()...Start a training job with hyperparameters.{"HyperParameters":{"batch-size":256,"learning-rate":0.0001,"communicator":"pure_nccl"}}Read additional information using environment variablesAn entry point often needs additional information not available inhyperparameters.
The SageMaker Training Toolkit writes this information as environment variables that are available from within the script.
For example, this training job includes the channelstrainingandtesting:fromsagemaker.pytorchimportPyTorchestimator=PyTorch(entry_point="train.py",...)estimator.fit({"training":"s3://bucket/path/to/training/data","testing":"s3://bucket/path/to/testing/data"})The environment variablesSM_CHANNEL_TRAININGandSM_CHANNEL_TESTINGprovide the paths to the channels:importargparseimportosif__name__=="__main__":parser=argparse.ArgumentParser()...# reads input channels training and testing from the environment variablesparser.add_argument("--training",type=str,default=os.environ["SM_CHANNEL_TRAINING"])parser.add_argument("--testing",type=str,default=os.environ["SM_CHANNEL_TESTING"])args=parser.parse_args()...When training starts, SageMaker Training Toolkit will print all available environment variables. Please see thereference on environment variablesfor a full list of provided environment variables.Get information about the container environmentTo get information about the container environment, initialize anEnvironmentobject.Environmentprovides access to aspects of the environment relevant to training jobs, including hyperparameters, system characteristics, filesystem locations, environment variables and configuration settings.
It is a read-only snapshot of the container environment during training, and it doesn't contain any form of state.fromsagemaker_trainingimportenvironmentenv=environment.Environment()# get the path of the channel "training" from the `inputdataconfig.json` filetraining_dir=env.channel_input_dirs["training"]# get a the hyperparameter "training_data_file" from `hyperparameters.json` filefile_name=env.hyperparameters["training_data_file"]# get the folder where the model should be savedmodel_dir=env.model_dir# train the modeldata=np.load(os.path.join(training_dir,file_name))x_train,y_train=data["features"],keras.utils.to_categorical(data["labels"])model=ResNet50(weights="imagenet")...model.fit(x_train,y_train)#save the model to the model_dir at the end of trainingmodel.save(os.path.join(model_dir,"saved_model"))Execute the entry pointTo execute the entry point, callentry_point.run().fromsagemaker_trainingimportentry_point,environmentenv=environment.Environment()# read hyperparameters as script argumentsargs=env.to_cmd_args()# get the environment variablesenv_vars=env.to_env_vars()# execute the entry pointentry_point.run(uri=env.module_dir,user_entry_point=env.user_entry_point,args=args,env_vars=env_vars)If the entry point execution fails,trainer.train()will write the error message to/opt/ml/output/failure. Otherwise, it will write to the file/opt/ml/success.:scroll: LicenseThis library is licensed under theApache 2.0 License.
For more details, please take a look at theLICENSEfile.:handshake: ContributingContributions are welcome!
Please read ourcontributing guidelinesif you'd like to open an issue or submit a pull request.