torchvision VOC dataset download 문제

torchvision.datasets.VOCDetection를 상속받아 클래스를 만든 후 생성 시 download 매개변수를 True를 주면

VOC 데이터셋을 자동으로 다운로드 받을 수 있다.

 

상세한 매개변수는 voc.py의 class _VOCBase를 참조하면 된다.

기본값은 아래와 같다.

class _VOCBase(VisionDataset):
    _SPLITS_DIR: str
    _TARGET_DIR: str
    _TARGET_FILE_EXT: str

    def __init__(
        self,
        root: str,
        year: str = "2012",
        image_set: str = "train",
        download: bool = False,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        transforms: Optional[Callable] = None,
    ):
...

그러나 torchvision 0.9.0의 voc.py에는 year=2007일 때, test만 다운로드 받아지는 문제가 있는데 코드가 잘못 작성되어 있다.

Link: torchvision 0.9.0 voc

 

87번째 라인부터 보면 다음과 같이 되어있는데

  valid_image_sets = ["train", "trainval", "val"]
  if year == "2007":
  	valid_image_sets.append("test")
  	key = "2007-test"
  else:
  	key = year
  self.image_set = verify_str_arg(image_set, "image_set", valid_image_sets)
  dataset_year_dict = DATASET_YEAR_DICT[key]

key에 2007-test를 할당하여 무조건 test 데이터셋만 받아진다.

 

그 외의 버전은

valid_image_sets = ["train", "trainval", "val"]
if year == "2007":
	valid_image_sets.append("test")
self.image_set = verify_str_arg(image_set, "image_set", valid_image_sets)

key = "2007-test" if year == "2007" and image_set == "test" else year
dataset_year_dict = DATASET_YEAR_DICT[key]

조건연산자를 이용하여 제대로 작성되어 있다.

이 부분의 코드를 이렇게 수정을 하거나 직접 URL을 통하여 다운로드 받거나

혹은 다른 버전의 torchvision을 다운로드 받으면 된다.