preprocess-bench.py 2.29 KB
Newer Older
soumith's avatar
soumith committed
1
2
3
import argparse
import os
from timeit import default_timer as timer
4

soumith's avatar
soumith committed
5
6
import torch
import torch.utils.data
7
import torchvision
soumith's avatar
soumith committed
8
import torchvision.datasets as datasets
9
10
import torchvision.transforms as transforms
from torch.utils.model_zoo import tqdm
soumith's avatar
soumith committed
11
12


13
14
15
16
17
18
19
20
21
parser = argparse.ArgumentParser(description="PyTorch ImageNet Training")
parser.add_argument("--data", metavar="PATH", required=True, help="path to dataset")
parser.add_argument(
    "--nThreads", "-j", default=2, type=int, metavar="N", help="number of data loading threads (default: 2)"
)
parser.add_argument(
    "--batchSize", "-b", default=256, type=int, metavar="N", help="mini-batch size (1 = pure stochastic) Default: 256"
)
parser.add_argument("--accimage", action="store_true", help="use accimage")
soumith's avatar
soumith committed
22
23
24
25
26


if __name__ == "__main__":
    args = parser.parse_args()

27
    if args.accimage:
28
        torchvision.set_image_backend("accimage")
29
    print(f"Using {torchvision.get_image_backend()}")
30

soumith's avatar
soumith committed
31
    # Data loading code
32
33
34
35
36
37
38
39
40
41
42
    transform = transforms.Compose(
        [
            transforms.RandomSizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )

    traindir = os.path.join(args.data, "train")
    valdir = os.path.join(args.data, "val")
soumith's avatar
soumith committed
43
44
45
    train = datasets.ImageFolder(traindir, transform)
    val = datasets.ImageFolder(valdir, transform)
    train_loader = torch.utils.data.DataLoader(
46
47
        train, batch_size=args.batchSize, shuffle=True, num_workers=args.nThreads
    )
soumith's avatar
soumith committed
48
49
50
    train_iter = iter(train_loader)

    start_time = timer()
51
    batch_count = 20 * args.nThreads
Francisco Massa's avatar
Francisco Massa committed
52
53
54
55
    with tqdm(total=batch_count) as pbar:
        for _ in tqdm(range(batch_count)):
            pbar.update(1)
            batch = next(train_iter)
soumith's avatar
soumith committed
56
    end_time = timer()
57
58
59
60
61
62
63
64
65
    print(
        "Performance: {dataset:.0f} minutes/dataset, {batch:.1f} ms/batch,"
        " {image:.2f} ms/image {rate:.0f} images/sec".format(
            dataset=(end_time - start_time) * (float(len(train_loader)) / batch_count / 60.0),
            batch=(end_time - start_time) / float(batch_count) * 1.0e3,
            image=(end_time - start_time) / (batch_count * args.batchSize) * 1.0e3,
            rate=(batch_count * args.batchSize) / (end_time - start_time),
        )
    )