Merge branch 'main' into imagenet

Author: NicolasHug

packaging/build_conda.sh

@@ -13,8 +13,4 @@ setup_visual_studio_constraint
 setup_junit_results_folder
 export CUDATOOLKIT_CHANNEL="nvidia"
 
-if [[ "$PYTHON_VERSION" == "3.11" ]]; then
-  export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c malfet"
-fi
-
 conda build -c $CUDATOOLKIT_CHANNEL $CONDA_CHANNEL_FLAGS --no-anaconda-upload --no-test --python "$PYTHON_VERSION" packaging/torchvision

torchvision/_meta_registrations.py

@@ -1,3 +1,5 @@
+import functools
+
 import torch
 import torch.library
 
@@ -6,20 +8,22 @@
 
 from torch._prims_common import check
 
-_meta_lib = torch.library.Library("torchvision", "IMPL", "Meta")
 
-vision = torch.ops.torchvision
+@functools.lru_cache(None)
+def get_meta_lib():
+    return torch.library.Library("torchvision", "IMPL", "Meta")
 
 
-def register_meta(op):
+def register_meta(op_name, overload_name="default"):
     def wrapper(fn):
-        _meta_lib.impl(op, fn)
+        if torchvision.extension._has_ops():
+            get_meta_lib().impl(getattr(getattr(torch.ops.torchvision, op_name), overload_name), fn)
         return fn
 
     return wrapper
 
 
-@register_meta(vision.roi_align.default)
+@register_meta("roi_align")
 def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned):
     check(rois.size(1) == 5, lambda: "rois must have shape as Tensor[K, 5]")
     check(
@@ -34,7 +38,7 @@ def meta_roi_align(input, rois, spatial_scale, pooled_height, pooled_width, samp
     return input.new_empty((num_rois, channels, pooled_height, pooled_width))
 
 
-@register_meta(vision._roi_align_backward.default)
+@register_meta("_roi_align_backward")
 def meta_roi_align_backward(
     grad, rois, spatial_scale, pooled_height, pooled_width, batch_size, channels, height, width, sampling_ratio, aligned
 ):

torchvision/datapoints/_dataset_wrapper.py

@@ -38,8 +38,8 @@ def wrap_dataset_for_transforms_v2(dataset, target_keys=None):
         * :class:`~torchvision.datasets.CocoDetection`: Instead of returning the target as list of dicts, the wrapper
           returns a dict of lists. In addition, the key-value-pairs ``"boxes"`` (in ``XYXY`` coordinate format),
           ``"masks"`` and ``"labels"`` are added and wrap the data in the corresponding ``torchvision.datapoints``.
-          The original keys are preserved. If ``target_keys`` is ommitted, returns only the values for the ``"boxes"``
-          and ``"labels"``.
+          The original keys are preserved. If ``target_keys`` is ommitted, returns only the values for the
+          ``"image_id"``, ``"boxes"``, and ``"labels"``.
         * :class:`~torchvision.datasets.VOCDetection`: The key-value-pairs ``"boxes"`` and ``"labels"`` are added to
           the target and wrap the data in the corresponding ``torchvision.datapoints``. The original keys are
           preserved. If ``target_keys`` is ommitted, returns only the values for the ``"boxes"`` and ``"labels"``.
@@ -328,7 +328,7 @@ def coco_dectection_wrapper_factory(dataset, target_keys):
             "masks",
             "labels",
         },
-        default={"boxes", "labels"},
+        default={"image_id", "boxes", "labels"},
     )
 
     def segmentation_to_mask(segmentation, *, spatial_size):