Training an object detection model with the TensorFlow Object Detection API and Amazon SageMaker can be a powerful way to combine deep learning with scalable cloud infrastructure. However, setting up the environment can sometimes be challenging, especially when working with older tutorials and newer versions of the AWS SDK.
In this post, I will describe the main issues I encountered while working on a TensorFlow 2 Object Detection project using the Waymo Open Dataset, AWS SageMaker, Docker, and Amazon ECR, together with the solutions that allowed the training workflow to work correctly.
Project Overview
The project uses the TensorFlow Object Detection API to train and evaluate different object detection architectures on the Waymo Open Dataset.
The dataset was already converted into TFRecord format and stored in Amazon S3. The images have a resolution of 640 × 640 pixels.
The main workflow is:
Store the training and validation data in Amazon S3.
Build a Docker container containing the TensorFlow Object Detection API.
Push the Docker image to Amazon Elastic Container Registry (ECR).
Download a pretrained model from the TensorFlow Model Zoo.
Configure the training pipeline.
Launch a SageMaker training job.
Evaluate the model and compare different experiments.
The architectures considered include SSD MobileNet, SSD ResNet50, Faster R-CNN, EfficientDet, and Faster R-CNN ResNet152.
Problem 1: sagemaker.estimator Cannot Be Imported
The first error I encountered was:
ModuleNotFoundError: No module named 'sagemaker.estimator'
The problematic import was:
import sagemaker
from sagemaker.estimator import Estimator
from framework import CustomFramework
The reason was a compatibility issue between the notebook and the installed SageMaker Python SDK.
The original project was designed around SageMaker SDK version 2, while newer environments may install version 3 by default. The SageMaker SDK v3 introduced significant API changes, so older code relying on classes such as Estimator may no longer work.
Solution
Instead of installing the newest version of SageMaker, I installed a version from the 2.x series:
%pip install "sagemaker<3" tensorflow_io
After restarting the notebook kernel, the following imports worked again:
import sagemaker
from sagemaker.estimator import Estimator
from framework import CustomFramework
This produced a deprecation warning explaining that SageMaker SDK v2 is no longer the actively developed version.
For this particular project, however, using SDK v2 is appropriate because the provided CustomFramework implementation was designed around the older API.
Problem 2: SageMaker Could Not Find the Docker Image
After solving the SDK issue, the training job failed with another error:
An error occurred (ValidationException) when calling the CreateTrainingJob operation:
Cannot find the requested image:
166664655187.dkr.ecr.us-east-1.amazonaws.com/tf2-object-detection:20260817082045
At first, this looked like a SageMaker problem, but the important part of the message was the ECR image name and tag.
SageMaker was trying to download an image from Amazon ECR, but that exact image did not exist.
The project uses a script similar to:
./docker/build_and_push.sh tf2-object-detection
to build the Docker image and push it to ECR.
I checked the repository with:
aws ecr describe-images \
--repository-name tf2-object-detection \
--region us-east-1
This helped determine whether the timestamped image tag actually existed.
Problem 3: Docker Was Not Authorized to Push to ECR
The final error revealed the real cause:
denied: User:
arn:aws:sts::166664655187:assumed-role/AmazonSageMaker-ExecutionRole-20260817T095324/SageMaker
is not authorized to perform:
ecr:InitiateLayerUpload
This was the key discovery.
The SageMaker execution role did not have enough permissions to upload Docker image layers to Amazon ECR.
The build script was therefore unable to push the image successfully. However, it still wrote an image URI into:
docker/ecr_image_fullname.txt
This explained the previous error: SageMaker received an image URI, but the corresponding image had never actually been uploaded to ECR.
Granting ECR Permissions
The SageMaker execution role needs permissions that allow it to interact with ECR.
For a development or educational project, one straightforward solution is to attach the AWS-managed:
AmazonEC2ContainerRegistryPowerUser
policy to the SageMaker execution role.
The relevant role in my environment was:
AmazonSageMaker-ExecutionRole-20260817T095324
In the AWS IAM console, the process is:
Open IAM.
Open Roles.
Select the SageMaker execution role.
Choose Add permissions.
Select Attach policies.
Search for
AmazonEC2ContainerRegistryPowerUser.Attach the policy.
For production environments, it is generally better to use a more restrictive custom IAM policy that grants only the ECR actions and repository resources actually required.
Rebuilding and Pushing the Image
After updating the IAM permissions, I ran the Docker build and push script again:
./docker/build_and_push.sh tf2-object-detection
This time the Docker image could be uploaded to ECR successfully.
I then verified the repository:
aws ecr describe-images \
--repository-name tf2-object-detection \
--region us-east-1
The newly generated timestamp tag was now visible.
I also refreshed the container variable:
with open('docker/ecr_image_fullname.txt', 'r') as f:
container = f.read().strip()
print(container)
The value looked similar to:
166664655187.dkr.ecr.us-east-1.amazonaws.com/tf2-object-detection:20260817103000
The important point is that this exact tag must exist in ECR before starting the SageMaker training job.
Launching the Training Job
Once the Docker image was available in ECR, the SageMaker estimator could be created:
estimator = CustomFramework(
role=role,
image_uri=container,
entry_point='run_training.sh',
source_dir='source_dir/',
hyperparameters={
"model_dir": "/opt/training",
"pipeline_config_path": "pipeline.config",
"num_train_steps": "2000",
"sample_1_of_n_eval_examples": "1"
},
instance_count=1,
instance_type='ml.g5.xlarge',
tensorboard_output_config=tensorboard_output_config,
disable_profiler=True,
base_job_name='tf2-object-detection'
)
The training job can then be started with:
estimator.fit(inputs)
SageMaker launches the specified GPU instance, downloads the Docker image from ECR, accesses the training data from S3, and executes the training script inside the container.
Lessons Learned
This project highlighted several important lessons about machine learning workflows in the cloud.
1. Check software versions
Older machine learning tutorials can depend on APIs that have changed significantly.
Before troubleshooting the code itself, it is useful to check:
import sagemaker
print(sagemaker.__version__)
Using the SDK version expected by the project can save considerable debugging time.
2. An image URI does not guarantee that an image exists
A file such as:
ecr_image_fullname.txt
may contain a valid-looking ECR URI even if the Docker push failed.
Always verify the image directly:
aws ecr describe-images \
--repository-name tf2-object-detection \
--region us-east-1
3. IAM permissions are critical
AWS services frequently interact through IAM roles. A SageMaker execution role may have permission to run training jobs but still lack permission to push or pull Docker images from ECR.
The specific error message is often very useful because it identifies the missing action, such as:
ecr:InitiateLayerUpload
4. Verify the complete pipeline before starting an expensive training job
GPU instances such as ml.g5.xlarge can be expensive. It is better to verify the following before launching training:
The S3 training data is accessible.
The S3 validation data is accessible.
The Docker image builds correctly.
The Docker image is successfully pushed to ECR.
The ECR image tag exists.
The SageMaker execution role has the required permissions.
The
pipeline.configfile matches the selected model architecture.
This avoids wasting compute time on infrastructure problems.
Next Steps: Improving the Object Detection Model
Once the infrastructure is working, the next challenge is improving model performance.
The initial experiment can be used as a baseline. Further experiments can focus on:
Data augmentation.
Learning rate and optimizer configuration.
Number of training steps.
Batch size.
Different object detection architectures.
Different pretrained checkpoints.
Evaluation frequency.
Model-specific hyperparameters.
For example, TensorFlow's Object Detection API provides several augmentation techniques that can improve robustness to variations in lighting, scale, orientation, and image composition.
The experiments should be evaluated using metrics such as mean Average Precision (mAP), while TensorBoard can be used to analyze training and validation behavior.
The biggest challenge in this project was not necessarily the object detection model itself, but making the different components of the cloud training pipeline work together.
The final workflow can be summarized as:
Waymo Dataset
↓
Amazon S3
↓
TensorFlow Object Detection API
↓
Docker Container
↓
Amazon ECR
↓
Amazon SageMaker
↓
GPU Training
↓
Evaluation + TensorBoard
The two most important troubleshooting steps were using a compatible SageMaker SDK version and ensuring that the SageMaker execution role had the necessary ECR permissions.
Once these infrastructure issues were resolved, SageMaker could successfully use the custom TensorFlow Object Detection container and proceed with model training.
This experience also reinforced an important principle for cloud-based machine learning: before optimizing the model, make sure the entire data, container, permissions, and training pipeline is working reliably.