def test_imagenet(self): train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=256, shuffle=False) random.seed(0) tc_data = [] for i, data in enumerate(train_loader): tc_data.append(data) print("get", data[0].shape) if i == check_num_batch: break from jittor.dataset.dataset import ImageFolder import jittor.transform as transform dataset = ImageFolder(traindir).set_attrs(batch_size=256, shuffle=False) dataset.set_attrs(transform=transform.Compose([ transform.RandomCropAndResize(224), transform.RandomHorizontalFlip(), transform.ImageNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])) random.seed(0) for i, (images, labels) in enumerate(dataset): print("compare", i) assert np.allclose(images.numpy(), tc_data[i][0].numpy()) assert np.allclose(labels.numpy(), tc_data[i][1].numpy()) if i == check_num_batch: break
labels = [path.split('/')[-1].split('.')[0] for path in train_list] ## Split train_list, valid_list = train_test_split(train_list, test_size=0.2, stratify=labels, random_state=42) print(f"Train Data: {len(train_list)}") print(f"Validation Data: {len(valid_list)}") print(f"Test Data: {len(test_list)}") ## Image Augumentation transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomCropAndResize(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) ## Load Datasets class CatsDogsDataset(Dataset): def __init__(self, file_list, transform=None, batch_size=1, shuffle=False, num_workers=0): super(CatsDogsDataset, self).__init__(batch_size=batch_size, shuffle=shuffle,
parser.add_argument('--cuda', action='store_true', help='to use cuda') args = parser.parse_args() if not os.path.isdir(args.logdir): os.mkdir(args.logdir) logging.basicConfig(filename=os.path.join(args.logdir, 'log.txt'), level=logging.INFO) logging.info(sys.argv) jt.flags.use_cuda = 1 if args.cuda else 0 train_transform = transform.Compose([ transform.RandomHorizontalFlip(), transform.RandomCropAndResize(32, scale=(0.5, 1)), transform.ImageNormalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) ]) test_transform = transform.Compose([ transform.ImageNormalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) ]) trainloader = CIFAR10(train=True, shuffle=True, batch_size=64, transform=train_transform) testloader = CIFAR10(train=False, shuffle=False,