watermarking-cnn-classifiers/project_datasets.py at main · gandizm/watermarking-cnn-classifiers · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

from collections import namedtuple

import torch
import torchvision
import torchvision.transforms as transforms
from torch.autograd import Variable

import watermarking

_DatasetMeta = namedtuple("DatasetMeta", "name shape data_mean data_std "
                                         "transform_train transform_test "
                                         "data_dir batch_size num_train "
                                         "watermarked_dataset")


Cifar10Meta = _DatasetMeta(
    name='CIFAR10',
    shape=[3, 32, 32],
    data_mean=[0.4914, 0.4822, 0.4465],
    data_std=[0.2023, 0.1994, 0.2010],
    num_train=50000,
    batch_size=128,
    transform_train=transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]),
    transform_test=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]),
    data_dir='cifar10_pytorch',
    watermarked_dataset=watermarking.WatermarkedCIFAR10,
)

Cifar100Meta = _DatasetMeta(
    name='CIFAR100',
    shape=[3, 32, 32],
    data_mean=[0.4914, 0.4822, 0.4465],
    data_std=[0.2023, 0.1994, 0.2010],
    num_train=50000,
    batch_size=128,
    transform_train=transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]),
    transform_test=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ]),
    data_dir='cifar10_pytorch',
    watermarked_dataset=watermarking.WatermarkedCIFAR100,
)

MnistMeta = _DatasetMeta(
    name='MNIST',
    shape=[1, 28, 28],
    data_mean=[0.1307],
    data_std=[0.3081],
    num_train=60000,
    batch_size=128,
    transform_train=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ]),
    transform_test=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ]),
    data_dir='mnist_pytorch',
    watermarked_dataset=watermarking.WatermarkedMNIST,
)