- 文件夹顺序
from torchvision import datasets
data_dir=$DATA
train_datasets = datasets.ImageFolder(os.path.join(data_dir, 'train')
print(train_datasets.classes, train_datasets.class_to_idx)
['negative', 'others', 'positive'],{'negative': 0, 'others': 1, 'positive': 2}
def find_classes(dir):
classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
classes.sort()
class_to_idx = {classes[i]: i for i in range(len(classes))}
return classes, class_to_idx
- 直接获取文件夹顺序
ls $DATA
negative others positive
参考文献: