| 26 | |
| 27 | |
| 28 | def get_loader(image_size): |
| 29 | transform = transforms.Compose( |
| 30 | [ |
| 31 | transforms.Resize((image_size, image_size)), |
| 32 | transforms.ToTensor(), |
| 33 | transforms.RandomHorizontalFlip(p=0.5), |
| 34 | transforms.Normalize( |
| 35 | [0.5 for _ in range(config.CHANNELS_IMG)], |
| 36 | [0.5 for _ in range(config.CHANNELS_IMG)], |
| 37 | ), |
| 38 | ] |
| 39 | ) |
| 40 | #根据选择图片的大小来选择BatchSize |
| 41 | batch_size = config.BATCH_SIZES[int(log2(image_size / 4))] |
| 42 | #加载数据集 |
| 43 | dataset = datasets.ImageFolder(root=config.DATASET, transform=transform) |
| 44 | loader = DataLoader( |
| 45 | dataset, |
| 46 | batch_size=batch_size, |
| 47 | shuffle=True, |
| 48 | num_workers=config.NUM_WORKERS, |
| 49 | pin_memory=True, |
| 50 | ) |
| 51 | return loader, dataset |
| 52 | |
| 53 | |
| 54 | def train_fn( |