#!/usr/bin/env python# -- coding: utf-8 --# Python version: 3.6import copyimport torchimport numpy as npfrom options import args_parserfrom torch.utils.data import DataLoaderfrom torch.utils.data import Datasetfrom torchvision import datasets, transformsfrom sampling import visDrone_iid, visDrone_noniid, visDrone_noniid_unequalfrom util.dataloader import YoloDataset, yolo_dataset_collate# from sampling import cifar_iid, cifar_noniiddef get_dataset(args): """ Returns train and test datasets and a user group which is a dict where the keys are the user index and the values are the corresponding data for each of those users. """ if args.dataset=='visDrone': # data_dir = '../data/' # apply_transform = transforms.Compose( # [transforms.ToTensor(), # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) xmlfilepath = r'./data/train/annotations_voc' saveBasePath = r"./data/train/train_namelist/" annotation_path = '2007_train.txt' mosaic = False input_shape = (416, 416) val_split = 0.1 with open(annotation_path) as f: lines = f.readlines() np.random.seed(10101) np.random.shuffle(lines) np.random.seed(None) num_val = int(len(lines) * val_split) num_train = len(lines) - num_val Batch_size = 2 train_dataset = YoloDataset(lines[:num_train], (input_shape[0], input_shape[1]), mosaic=mosaic, is_train=True) test_dataset = YoloDataset(lines[num_train:], (input_shape[0], input_shape[1]), mosaic=False, is_train=False) gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate) gen_test = DataLoader(test_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True, drop_last=True, collate_fn=yolo_dataset_collate) return gen, gen_testargs = args_parser()if args.gpu: torch.cuda.set_device(args.gpu)device = 'cuda' if args.gpu else 'cpu'train_dataset, test_dataset = get_dataset(args)trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True)for batch_idx, (images, labels) in enumerate(trainloader): images, labels = images.to(device), labels.to(device)
为什么代码运行会出现TypeError: 'DataLoader' object is not subscriptable问题
