[WIP] Closing streams"

ghstack-source-id: 37ba421
Pull Request resolved: #6128

Author: VitalyFedyunin

torchvision/prototype/datasets/_builtin/caltech.py

@@ -107,7 +107,9 @@ def _prepare_sample(
         ann_path, ann_buffer = ann_data
 
         image = EncodedImage.from_file(image_buffer)
+        image_buffer.close()
         ann = read_mat(ann_buffer)
+        ann_buffer.close()
 
         return dict(
             label=Label.from_category(category, categories=self._categories),
@@ -186,10 +188,11 @@ def _is_not_rogue_file(self, data: Tuple[str, Any]) -> bool:
 
     def _prepare_sample(self, data: Tuple[str, BinaryIO]) -> Dict[str, Any]:
         path, buffer = data
-
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         return dict(
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
             label=Label(int(pathlib.Path(path).parent.name.split(".", 1)[0]) - 1, categories=self._categories),
         )
 

torchvision/prototype/datasets/_builtin/cifar.py

@@ -66,7 +66,9 @@ def _resources(self) -> List[OnlineResource]:
 
     def _unpickle(self, data: Tuple[str, io.BytesIO]) -> Dict[str, Any]:
         _, file = data
-        return cast(Dict[str, Any], pickle.load(file, encoding="latin1"))
+        result = pickle.load(file, encoding="latin1")
+        file.close()
+        return cast(Dict[str, Any], result)
 
     def _prepare_sample(self, data: Tuple[np.ndarray, int]) -> Dict[str, Any]:
         image_array, category_idx = data

torchvision/prototype/datasets/_builtin/coco.py

@@ -174,9 +174,11 @@ def _classify_meta(self, data: Tuple[str, Any]) -> Optional[int]:
 
     def _prepare_image(self, data: Tuple[str, BinaryIO]) -> Dict[str, Any]:
         path, buffer = data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
         return dict(
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _prepare_sample(
@@ -187,9 +189,11 @@ def _prepare_sample(
         anns, image_meta = ann_data
 
         sample = self._prepare_image(image_data)
+
         # this method is only called if we have annotations
         annotations = cast(str, self._annotations)
         sample.update(self._ANN_DECODERS[annotations](self, anns, image_meta))
+        image_data[1].close()
         return sample
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/imagenet.py

@@ -109,10 +109,11 @@ def _prepare_test_data(self, data: Tuple[str, BinaryIO]) -> Tuple[None, Tuple[st
         return None, data
 
     def _classifiy_devkit(self, data: Tuple[str, BinaryIO]) -> Optional[int]:
+        name, binary_io = data
         return {
             "meta.mat": ImageNetDemux.META,
             "ILSVRC2012_validation_ground_truth.txt": ImageNetDemux.LABEL,
-        }.get(pathlib.Path(data[0]).name)
+        }.get(pathlib.Path(name).name)
 
     # Although the WordNet IDs (wnids) are unique, the corresponding categories are not. For example, both n02012849
     # and n03126707 are labeled 'crane' while the first means the bird and the latter means the construction equipment
@@ -123,12 +124,14 @@ def _classifiy_devkit(self, data: Tuple[str, BinaryIO]) -> Optional[int]:
 
     def _extract_categories_and_wnids(self, data: Tuple[str, BinaryIO]) -> List[Tuple[str, str]]:
         synsets = read_mat(data[1], squeeze_me=True)["synsets"]
-        return [
+        results = [
             (self._WNID_MAP.get(wnid, category.split(",", 1)[0]), wnid)
             for _, wnid, category, _, num_children, *_ in synsets
             # if num_children > 0, we are looking at a superclass that has no direct instance
             if num_children == 0
         ]
+        data[1].close()
+        return results
 
     def _imagenet_label_to_wnid(self, imagenet_label: str, *, wnids: Tuple[str, ...]) -> str:
         return wnids[int(imagenet_label) - 1]
@@ -151,11 +154,13 @@ def _prepare_sample(
         data: Tuple[Optional[Tuple[Label, str]], Tuple[str, BinaryIO]],
     ) -> Dict[str, Any]:
         label_data, (path, buffer) = data
+        image = EncodedImage.from_file(buffer)
+        buffer.close()
 
         return dict(
             dict(zip(("label", "wnid"), label_data if label_data else (None, None))),
             path=path,
-            image=EncodedImage.from_file(buffer),
+            image=image,
         )
 
     def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str, Any]]:

torchvision/prototype/datasets/_builtin/mnist.py

@@ -37,6 +37,7 @@ def __init__(
 
     def __iter__(self) -> Iterator[torch.Tensor]:
         for _, file in self.datapipe:
+            print(file)
             read = functools.partial(fromfile, file, byte_order="big")
 
             magic = int(read(dtype=torch.int32, count=1))
@@ -57,6 +58,8 @@ def __iter__(self) -> Iterator[torch.Tensor]:
             for _ in range(stop - start):
                 yield read(dtype=dtype, count=count).reshape(shape)
 
+            file.close()
+
 
 class _MNISTBase(Dataset):
     _URL_BASE: Union[str, Sequence[str]]

torchvision/prototype/datasets/_builtin/sbd.py

@@ -28,6 +28,7 @@
 
 NAME = "sbd"
 
+from torchdata.datapipes.utils import StreamWrapper
 
 @register_info(NAME)
 def _info() -> Dict[str, Any]:
@@ -89,10 +90,12 @@ def _prepare_sample(self, data: Tuple[Tuple[Any, Tuple[str, BinaryIO]], Tuple[st
         ann_path, ann_buffer = ann_data
 
         anns = read_mat(ann_buffer, squeeze_me=True)["GTcls"]
-
+        ann_buffer.close()
+        image = EncodedImage.from_file(image_buffer)
+        image_buffer.close()
         return dict(
             image_path=image_path,
-            image=EncodedImage.from_file(image_buffer),
+            image=image,
             ann_path=ann_path,
             # the boundaries are stored in sparse CSC format, which is not supported by PyTorch
             boundaries=_Feature(np.stack([raw_boundary.toarray() for raw_boundary in anns["Boundaries"].item()])),
@@ -111,6 +114,8 @@ def _datapipe(self, resource_dps: List[IterDataPipe]) -> IterDataPipe[Dict[str,
             drop_none=True,
         )
         if self._split == "train_noval":
+            for i in split_dp:
+                StreamWrapper.cleanup_structure(i)
             split_dp = extra_split_dp
 
         split_dp = Filter(split_dp, path_comparator("name", f"{self._split}.txt"))

torchvision/prototype/datasets/_builtin/voc.py

@@ -101,7 +101,9 @@ def _classify_archive(self, data: Tuple[str, Any]) -> Optional[int]:
             return None
 
     def _parse_detection_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
-        return cast(Dict[str, Any], VOCDetection.parse_voc_xml(ElementTree.parse(buffer).getroot())["annotation"])
+        result = cast(Dict[str, Any], VOCDetection.parse_voc_xml(ElementTree.parse(buffer).getroot())["annotation"])
+        buffer.close()
+        return result
 
     def _prepare_detection_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
         anns = self._parse_detection_ann(buffer)
@@ -121,7 +123,9 @@ def _prepare_detection_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
         )
 
     def _prepare_segmentation_ann(self, buffer: BinaryIO) -> Dict[str, Any]:
-        return dict(segmentation=EncodedImage.from_file(buffer))
+        result = dict(segmentation=EncodedImage.from_file(buffer))
+        buffer.close()
+        return result
 
     def _prepare_sample(
         self,
@@ -132,10 +136,13 @@ def _prepare_sample(
         image_path, image_buffer = image_data
         ann_path, ann_buffer = ann_data
 
+        image = EncodedImage.from_file(image_buffer)
+        image_buffer.close()
+
         return dict(
             (self._prepare_detection_ann if self._task == "detection" else self._prepare_segmentation_ann)(ann_buffer),
             image_path=image_path,
-            image=EncodedImage.from_file(image_buffer),
+            image=image,
             ann_path=ann_path,
         )