Merge branch 'mvit' into mvitv1

Author: datumbox

.circleci/unittest/linux/scripts/install.sh

@@ -21,7 +21,12 @@ else
     fi
     echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="nvidia::cudatoolkit=${version}"
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        cuda_toolkit_pckg="cuda"
+    fi
+    cudatoolkit="nvidia::${cuda_toolkit_pckg}=${version}"
 fi
 
 case "$(uname -s)" in

.circleci/unittest/windows/scripts/install.sh

@@ -22,9 +22,15 @@ else
     elif [[ ${#CU_VERSION} -eq 5 ]]; then
         CUDA_VERSION="${CU_VERSION:2:2}.${CU_VERSION:4:1}"
     fi
+
+    cuda_toolkit_pckg="cudatoolkit"
+    if [[ "$CU_VERSION" == cu116 ]]; then
+        cuda_toolkit_pckg="cuda"
+    fi
+
     echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION"
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-    cudatoolkit="cudatoolkit=${version}"
+    cudatoolkit="${cuda_toolkit_pckg}=${version}"
 fi
 
 printf "Installing PyTorch with %s\n" "${cudatoolkit}"

packaging/build_conda.sh

@@ -11,11 +11,6 @@ setup_conda_pytorch_constraint
 setup_conda_cudatoolkit_constraint
 setup_visual_studio_constraint
 setup_junit_results_folder
-
-# nvidia channel included for cudatoolkit >= 11 however for 11.5 and 11.6 we use conda-forge
 export CUDATOOLKIT_CHANNEL="nvidia"
-if [[ "$CU_VERSION" == cu116 ]]; then
-    export CUDATOOLKIT_CHANNEL="conda-forge"
-fi
 
 conda build -c $CUDATOOLKIT_CHANNEL -c defaults $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchvision

packaging/pkg_helpers.bash

@@ -257,7 +257,7 @@ setup_conda_cudatoolkit_constraint() {
   else
     case "$CU_VERSION" in
       cu116)
-        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.6,<11.7 # [not osx]"
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cuda >=11.6,<11.7 # [not osx]"
         ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=11.3,<11.4 # [not osx]"
@@ -286,7 +286,7 @@ setup_conda_cudatoolkit_plain_constraint() {
   else
     case "$CU_VERSION" in
       cu116)
-        export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.6"
+        export CONDA_CUDATOOLKIT_CONSTRAINT="cuda=11.6"
         ;;
       cu113)
         export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit=11.3"

references/detection/train.py

@@ -35,6 +35,11 @@
 from transforms import SimpleCopyPaste
 
 
+def copypaste_collate_fn(batch):
+    copypaste = SimpleCopyPaste(blending=True, resize_interpolation=InterpolationMode.BILINEAR)
+    return copypaste(*utils.collate_fn(batch))
+
+
 def get_dataset(name, image_set, transform, data_path):
     paths = {"coco": (data_path, get_coco, 91), "coco_kp": (data_path, get_coco_kp, 2)}
     p, ds_fn, num_classes = paths[name]
@@ -194,11 +199,6 @@ def main(args):
         if args.data_augmentation != "lsj":
             raise RuntimeError("SimpleCopyPaste algorithm currently only supports the 'lsj' data augmentation policies")
 
-        copypaste = SimpleCopyPaste(resize_interpolation=InterpolationMode.BILINEAR, blending=True)
-
-        def copypaste_collate_fn(batch):
-            return copypaste(*utils.collate_fn(batch))
-
         train_collate_fn = copypaste_collate_fn
 
     data_loader = torch.utils.data.DataLoader(