File size: 1,083 Bytes
caac576 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
from sagemaker.huggingface import HuggingFace
ROLE = ?
# hyperparameters, which are passed into the training job
hyperparameters = {
'epochs': 1,
'per_device_train_batch_size': 32,
'do_train': True,
'model_name_or_path': 'distilbert-base-uncased',
'output_dir': '/opt/ml/checkpoints'
}
# create the Estimator
huggingface_estimator = HuggingFace(
entry_point='train.py',
source_dir='.',
instance_type='local', # 'ml.p3.2xlarge',
instance_count=1,
checkpoint_s3_uri=f's3://{sess.default_bucket()}/checkpoints',
use_spot_instances=True,
max_wait=3600, # This should be equal to or greater than max_run in seconds'
max_run=1000,
role=ROLE,
transformers_version='4.4',
pytorch_version='1.6',
py_version='py36',
hyperparameters=hyperparameters,
)
huggingface_estimator.fit(
{
'train': 's3://sagemaker-us-east-1-558105141721/samples/datasets/imdb/train',
'test': 's3://sagemaker-us-east-1-558105141721/samples/datasets/imdb/test'
}
)
|