HuggingFace

HuggingFace Estimator

class sagemaker.huggingface.estimator.HuggingFace(py_version, entry_point, transformers_version=None, tensorflow_version=None, pytorch_version=None, source_dir=None, hyperparameters=None, image_uri=None, distribution=None, **kwargs)

Bases: sagemaker.estimator.Framework

Handle training of custom HuggingFace code.

This Estimator executes a HuggingFace script in a managed execution environment.

The managed HuggingFace environment is an Amazon-built Docker container that executes functions defined in the supplied entry_point Python script within a SageMaker Training Job.

Training is started by calling fit() on this Estimator.

Parameters
  • py_version (str) – Python version you want to use for executing your model training code. Defaults to None. Required unless image_uri is provided. If using PyTorch, the current supported version is py36. If using TensorFlow, the current supported version is py37.

  • entry_point (str) – Path (absolute or relative) to the Python source file which should be executed as the entry point to training. If source_dir is specified, then entry_point must point to a file located at the root of source_dir.

  • transformers_version (str) – Transformers version you want to use for executing your model training code. Defaults to None. Required unless image_uri is provided. The current supported version is 4.6.1.

  • tensorflow_version (str) – TensorFlow version you want to use for executing your model training code. Defaults to None. Required unless pytorch_version is provided. The current supported version is 1.6.0.

  • pytorch_version (str) – PyTorch version you want to use for executing your model training code. Defaults to None. Required unless tensorflow_version is provided. The current supported version is 2.4.1.

  • source_dir (str) – Path (absolute, relative or an S3 URI) to a directory with any other training source code dependencies aside from the entry point file (default: None). If source_dir is an S3 URI, it must point to a tar.gz file. Structure within this directory are preserved when training on Amazon SageMaker.

  • hyperparameters (dict) – Hyperparameters that will be used for training (default: None). The hyperparameters are made accessible as a dict[str, str] to the training code on SageMaker. For convenience, this accepts other types for keys and values, but str() will be called to convert them before training.

  • image_uri (str) –

    If specified, the estimator will use this image for training and hosting, instead of selecting the appropriate SageMaker official image based on framework_version and py_version. It can be an ECR url or dockerhub image and tag. .. rubric:: Examples

    • 123412341234.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0

    • custom-image:latest

    If framework_version or py_version are None, then image_uri is required. If also None, then a ValueError will be raised.

  • distribution (dict) –

    A dictionary with information on how to run distributed training (default: None). Currently, the following are supported: distributed training with parameter servers, SageMaker Distributed (SMD) Data and Model Parallelism, and MPI. SMD Model Parallelism can only be used with MPI. To enable parameter server use the following setup:

    {
        "parameter_server": {
            "enabled": True
        }
    }
    

    To enable MPI:

    {
        "mpi": {
            "enabled": True
        }
    }
    

    To enable SMDistributed Data Parallel or Model Parallel:

    {
        "smdistributed": {
            "dataparallel": {
                "enabled": True
            },
            "modelparallel": {
                "enabled": True,
                "parameters": {}
            }
        }
    }
    

  • **kwargs – Additional kwargs passed to the Framework constructor.

Tip

You can find additional parameters for initializing this class at Framework and EstimatorBase.

hyperparameters()

Return hyperparameters used by your custom PyTorch code during model training.

create_model(model_server_workers=None, role=None, vpc_config_override='VPC_CONFIG_DEFAULT', entry_point=None, source_dir=None, dependencies=None, **kwargs)

Create a SageMaker HuggingFaceModel object that can be deployed to an Endpoint.

Parameters
  • model_server_workers (int) – Optional. The number of worker processes used by the inference server. If None, server will use one worker per vCPU.

  • role (str) – The ExecutionRoleArn IAM Role ARN for the Model, which is also used during transform jobs. If not specified, the role from the Estimator will be used.

  • vpc_config_override (dict[str, list[str]]) – Optional override for VpcConfig set on the model. Default: use subnets and security groups from this Estimator. * ‘Subnets’ (list[str]): List of subnet ids. * ‘SecurityGroupIds’ (list[str]): List of security group ids.

  • entry_point (str) – Path (absolute or relative) to the local Python source file which should be executed as the entry point to training. If source_dir is specified, then entry_point must point to a file located at the root of source_dir. Defaults to None.

  • source_dir (str) – Path (absolute or relative) to a directory with any other serving source code dependencies aside from the entry point file. If not specified, the model source directory from training is used.

  • dependencies (list[str]) – A list of paths to directories (absolute or relative) with any additional libraries that will be exported to the container. If not specified, the dependencies from training are used. This is not supported with “local code” in Local Mode.

  • **kwargs – Additional kwargs passed to the HuggingFaceModel constructor.

Returns

A SageMaker HuggingFaceModel object. See HuggingFaceModel() for full details.

Return type

sagemaker.huggingface.model.HuggingFaceModel