feat: support flatten and reshape via shuffle_layer (#2354) #205
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: 'Torch-TensorRT Docker Build' | |
| # Apply workflow only to main branch | |
| on: | |
| push: | |
| branches: | |
| - main | |
| - nightly | |
| # If pushes to main are made in rapid succession, | |
| # cancel existing docker builds and use newer commits | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.ref_name }} | |
| cancel-in-progress: true | |
| jobs: | |
| build: | |
| runs-on: linux.2xlarge | |
| # Define key environment variables | |
| # Container name is of the form torch_tensorrt:<branch_name> | |
| env: | |
| DOCKER_REGISTRY: ghcr.io/pytorch/tensorrt | |
| CONTAINER_NAME: torch_tensorrt:${{ github.ref_name }} | |
| steps: | |
| - name: Checkout repository | |
| uses: actions/checkout@v3 | |
| - name: Log in to the Container registry | |
| uses: docker/login-action@v2 | |
| with: | |
| registry: ${{ env.DOCKER_REGISTRY }} | |
| username: ${{ github.actor }} | |
| password: ${{ secrets.GITHUB_TOKEN }} | |
| # Automatically detect TensorRT and cuDNN default versions for Torch-TRT build | |
| - name: Build Docker image | |
| env: | |
| DOCKER_TAG: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }} | |
| run: | | |
| python3 -m pip install pyyaml | |
| TRT_VERSION=$(python3 -c "import versions; versions.tensorrt_version()") | |
| echo "TRT VERSION = ${TRT_VERSION}" | |
| CUDNN_VERSION=$(python3 -c "import versions; versions.cudnn_version()") | |
| echo "CUDNN VERSION = ${CUDNN_VERSION}" | |
| DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=$TRT_VERSION --build-arg CUDNN_VERSION=$CUDNN_VERSION -f docker/Dockerfile --tag $DOCKER_TAG . | |
| - name: Push Docker image | |
| env: | |
| DOCKER_URL: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }} | |
| run: docker push $DOCKER_URL | |
| # Clean up all untagged containers in registry | |
| - name: Container Registry Cleanup | |
| uses: actions/delete-package-versions@v4 | |
| with: | |
| package-name: "tensorrt/torch_tensorrt" | |
| package-type: container | |
| min-versions-to-keep: 0 | |
| delete-only-untagged-versions: True |