def qmnist():
return itertools.chain(*[
collect_download_configs(
lambda: datasets.QMNIST(ROOT, what=what, download=True),
name=f"QMNIST, {what}",
file="mnist",
) for what in ("train", "test", "nist")
])
def __init__(self, transform=lambda x: x, what='train'):
super().__init__()
self.transform = transform
self.qmnist = datasets.QMNIST(root='resources/qmnist/', what=what, download=True)
self.printed = []
if os.path.exists('resources/csv/printed_digits.csv'):
with open('resources/csv/printed_digits.csv', 'r') as f:
reader = csv.reader(f, delimiter=',')
for row in reader:
self.printed.append((
base64ToPIL(row[2]).convert('L'),
int(row[1])
))
self.printed = [
x
for i, x in enumerate(self.printed)
if (i % 5 == 0) == (what == 'test')
]
def get_datasets(dataset_name,
*,
dataroot='/opt/ml/input/data',
image_size,
classes=None):
logger.info(f"dataname: {dataset_name}, dataroot: {dataroot}, " +
f"image_size: {image_size}")
if dataset_name in ['imagenet', 'folder', 'lfw']:
# folder dataset
dataset = dset.ImageFolder(root=dataroot,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
elif dataset_name == 'lsun':
classes = [c + '_train' for c in args.classes.split(',')]
dataset = dset.LSUN(root=dataroot,
classes=classes,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
nc = 3
elif dataset_name == 'cifar10':
dataset = dset.CIFAR10(root=dataroot,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
nc = 3
elif dataset_name == 'mnist':
dataset = dset.MNIST(root=dataroot,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
nc = 1
elif dataset_name == 'qmnist':
dataset = dset.QMNIST(root=dataroot,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
nc = 1
elif dataset_name == 'fashion-mnist':
dataset = dset.FashionMNIST(root=dataroot,
transform=transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
nc = 1
elif dataset_name == 'fake':
dataset = dset.FakeData(image_size=(3, image_size, image_size),
transform=transforms.ToTensor())
nc = 3
return dataset, nc
elif args.dataset == 'mnist':
dataset = dset.MNIST(root=args.dataroot,
download=True,
transform=transforms.Compose([
transforms.Resize(args.imageSize),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
nc = 1
elif args.dataset == 'qmnist':
dataset = dset.QMNIST(root=args.dataroot,
download=True,
transform=transforms.Compose([
transforms.Resize(args.imageSize),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
]))
nc = 1
elif args.dataset == 'fake':
dataset = dset.FakeData(image_size=(3, args.imageSize, args.imageSize),
transform=transforms.ToTensor())
nc = 3
num_workers = int(args.workers)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
def load_datasets(name, root, batch_size):
if name == "mnist":
train_dataset = datasets.MNIST(root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(28),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.MNIST(root=root,
download=True,
train=False,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
elif name == "fmnist":
train_dataset = datasets.FashionMNIST(
root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(28),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.FashionMNIST(root=root,
download=True,
train=False,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
transforms.Normalize([0.5],
[0.5]),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
elif name == "kmnist":
train_dataset = datasets.KMNIST(root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(28),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.KMNIST(root=root,
download=True,
train=False,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
elif name == "qmnist":
train_dataset = datasets.QMNIST(root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(28),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.QMNIST(root=root,
download=True,
what="test50k",
train=False,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
elif name == "cifar10":
train_dataset = datasets.CIFAR10(root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.CIFAR10(root=root,
download=True,
train=False,
transform=transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(
(0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
elif name == "cifar100":
train_dataset = datasets.CIFAR100(root=root,
download=True,
train=True,
transform=transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(
(0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
test_dataset = datasets.CIFAR100(root=root,
download=True,
train=False,
transform=transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(
(0.5, 0.5, 0.5),
(0.5, 0.5, 0.5)),
]))
test_dataloader = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
return train_dataloader, test_dataloader
def get_loader(args):
if args.exname == "AFS":
# Load train and test data directly.
if args.dataset == "MNIST":
train_dataset = datasets.MNIST(
root=args.data_root, train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(
root=args.data_root, train=False, transform=transforms.ToTensor())
elif args.dataset == "SVHN":
train_dataset = datasets.SVHN(
root=args.data_root, split="train", transform=transforms.Compose(_SVHN_TRAIN_TRANSFORMS), target_transform=transforms.Lambda(lambda y: y % 10), download=True
)
test_dataset = datasets.SVHN(root=args.data_root, split="test", transform=transforms.Compose(_SVHN_TEST_TRANSFORMS),
target_transform=transforms.Lambda(lambda y: y % 10), download=True)
elif args.dataset == "EMNIST":
train_dataset = datasets.EMNIST(
root=args.data_root, split="digits", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(
root=args.data_root, split="digits", train=False, transform=transforms.ToTensor(), download=True)
elif args.dataset == "KMNIST":
train_dataset = datasets.KMNIST(
root=args.data_root, train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.KMNIST(
root=args.data_root, train=False, transform=transforms.ToTensor(), download=True)
elif args.dataset == "QMNIST":
train_dataset = datasets.QMNIST(
root=args.data_root, what="train", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.QMNIST(
root=args.data_root, what="test", train=False, transform=transforms.ToTensor(), download=True)
elif args.dataset == "FashionMNIST":
train_dataset = datasets.FashionMNIST(
root=args.data_root, train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.FashionMNIST(
root=args.data_root, train=False, transform=transforms.ToTensor(), download=True)
else:
raise NotImplementedError
train_dataloader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=False, num_workers=args.num_workers, pin_memory=True)
test_dataloader = torch.utils.data.DataLoader(
test_dataset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.num_workers, pin_memory=True)
return train_dataloader, test_dataloader
elif args.exname == "TransferLearning":
# Load train dataset and test dataset for pretrain and finetune.
if args.dataset == "MNIST" and args.dataset_aux == "SVHN":
train_dataset = datasets.MNIST(
root=args.data_root, train=True, transform=transforms.Compose(_MNIST_COLORIZED_TRAIN_TRANSFORMS), download=True)
test_dataset = datasets.MNIST(
root=args.data_root, train=False, transform=transforms.Compose(_MNIST_COLORIZED_TEST_TRANSFORMS), download=True)
train_dataset_aux = datasets.SVHN(
root=args.data_root, split="train", transform=transforms.Compose(_SVHN_TRAIN_TRANSFORMS), target_transform=transforms.Lambda(lambda y: y % 10), download=True)
test_dataset_aux = datasets.SVHN(root=args.data_root, split="test", transform=transforms.Compose(
_SVHN_TEST_TRANSFORMS), target_transform=transforms.Lambda(lambda y: y % 10), download=True)
elif args.dataset == "SVHN" and args.dataset_aux == "MNIST":
train_dataset = datasets.SVHN(
root=args.data_root, split="train", transform=transforms.Compose(_SVHN_TRAIN_TRANSFORMS), target_transform=transforms.Lambda(lambda y: y % 10), download=True)
test_dataset = datasets.SVHN(root=args.data_root, split="test", transform=transforms.Compose(
_SVHN_TEST_TRANSFORMS), target_transform=transforms.Lambda(lambda y: y % 10), download=True)
train_dataset_aux = datasets.MNIST(
root=args.data_root, train=True, transform=transforms.Compose(_MNIST_COLORIZED_TRAIN_TRANSFORMS), download=True)
test_dataset_aux = datasets.MNIST(
root=args.data_root, train=False, transform=transforms.Compose(_MNIST_COLORIZED_TEST_TRANSFORMS), download=True)
elif args.dataset == "MNIST" and args.dataset_aux == "QMNIST":
train_dataset = datasets.MNIST(
root=args.data_root, train=True, transform=transforms.Compose(_MNIST_COLORIZED_TRAIN_TRANSFORMS), download=True)
test_dataset = datasets.MNIST(
root=args.data_root, train=False, transform=transforms.Compose(_MNIST_COLORIZED_TEST_TRANSFORMS), download=True)
train_dataset_aux = datasets.QMNIST(
root=args.data_root, what="train", train=True, transform=transforms.ToTensor(), download=True)
test_dataset_aux = datasets.QMNIST(
root=args.data_root, what="test", train=False, transform=transforms.ToTensor(), download=True)
elif args.dataset == "QMNIST" and args.dataset == "MNIST":
train_dataset = datasets.QMNIST(
root=args.data_root, what="train", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.QMNIST(
root=args.data_root, what="test", train=False, transform=transforms.ToTensor(), download=True)
train_dataset_aux = datasets.MNIST(
root=args.data_root, train=True, transform=transforms.Compose(_MNIST_COLORIZED_TRAIN_TRANSFORMS), download=True)
test_dataset_aux = datasets.MNIST(
root=args.data_root, train=False, transform=transforms.Compose(_MNIST_COLORIZED_TEST_TRANSFORMS), download=True)
else:
raise NotImplementedError
train_dataloader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=False, num_workers=args.num_workers, pin_memory=True)
test_dataloader = torch.utils.data.DataLoader(
test_dataset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.num_workers, pin_memory=True)
train_dataloader_aux = torch.utils.data.DataLoader(
train_dataset_aux, batch_size=args.batch_size, shuffle=True, drop_last=False, num_workers=args.num_workers, pin_memory=True)
test_dataloader_aux = torch.utils.data.DataLoader(
test_dataset_aux, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.num_workers, pin_memory=True)
return train_dataloader, test_dataloader, train_dataloader_aux, test_dataloader_aux
def preProcessData(p_Q,
args,
p_P=[.1, .1, .1, .1, .1, .1, .1, .1, .1, .1],
alpha=0.01,
img_size=32,
batch_size=128,
num_train_samples=30000,
num_test_samples=10000,
tweak_train=False,
tweak_type=0):
# data_loader
if args.dataset == "mnist":
args.img_size = 32
# img_shape= [1,args.img_size,args.img_size]
transform = transforms.Compose([
transforms.Resize(args.img_size),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, ))
])
train_loader = torch.utils.data.DataLoader(datasets.QMNIST(
'data', what='train', download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.QMNIST(
'data', what='test', download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
elif args.dataset == "cifar10":
transform = transforms.Compose([
transforms.Resize(args.img_size),
# transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
'data', train=True, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
'data', train=False, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
elif args.dataset == "cifar100":
transform = transforms.Compose([
transforms.Resize(args.img_size),
# transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.5070751592371323, 0.48654887331495095, 0.4409178433670343),
(0.2673342858792401, 0.2564384629170883, 0.27615047132568404)),
])
train_loader = torch.utils.data.DataLoader(datasets.CIFAR100(
'data', train=True, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.CIFAR100(
'data', train=False, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
elif args.dataset == "f-m":
transform = transforms.Compose([
transforms.Resize(args.img_size),
# transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
])
train_loader = torch.utils.data.DataLoader(datasets.FashionMNIST(
'data', train=True, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.FashionMNIST(
'data', train=False, download=True, transform=transform),
batch_size=batch_size,
shuffle=True)
for i, (x_, y_) in enumerate(train_loader):
if i == 0:
x_all, y_all = x_, y_
else:
x_all = np.concatenate((x_all, x_), axis=0)
y_all = np.concatenate((y_all, y_), axis=0)
for i, (x_, y_) in enumerate(test_loader):
if i == 0:
x_test, y_test = x_, y_
else:
x_test = np.concatenate((x_test, x_), axis=0)
y_test = np.concatenate((y_test, y_), axis=0)
num = 2
n = x_all.shape[0]
x_train, y_train = x_all, y_all
x_val, y_val = x_all, y_all
if tweak_train:
x_train, y_train = tweak_dist(x_train, y_train, args.num_class,
num_train_samples, p_P)
x_val, y_val = tweak_dist(x_val, y_val, args.num_class,
num_train_samples, p_P)
else:
x_test, y_test = tweak_dist(x_test, y_test, args.num_class,
num_test_samples, p_Q)
return x_train, y_train, x_val, y_val, x_test, y_test
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch QMNIST Example')
parser.add_argument('--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size',
type=int,
default=1000,
metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs',
type=int,
default=14,
metavar='N',
help='number of epochs to train (default: 14)')
parser.add_argument('--lr',
type=float,
default=1.0,
metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma',
type=float,
default=0.7,
metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument(
'--log-interval',
type=int,
default=10,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model',
action='store_true',
default=False,
help='For Saving the current Model')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader = torch.utils.data.DataLoader(datasets.QMNIST(
'data', train=True, download=True, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(datasets.QMNIST(
'data', train=False, download=True, transform=transforms.ToTensor()),
batch_size=args.test_batch_size,
shuffle=False,
**kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(args, model, device, test_loader)
scheduler.step()
if args.save_model:
torch.save(model.state_dict(), "qmnist_cnn.pt")
