From fa95c8d2d3663fc3f431beb52e792e5e469db449 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Mon, 3 Dec 2018 10:47:25 -0500 Subject: [PATCH] Cleaning up. --- mine_mnist.py | 84 ++++++++++++++++++++++++++++----------------------- 1 file changed, 46 insertions(+), 38 deletions(-) diff --git a/mine_mnist.py b/mine_mnist.py index 412c624..06458b5 100755 --- a/mine_mnist.py +++ b/mine_mnist.py @@ -1,11 +1,10 @@ #!/usr/bin/env python -import argparse +import argparse, math, sys -import math, sys, torch, torchvision +import torch, torchvision from torch import nn -from torch.nn import functional as F ###################################################################### @@ -28,27 +27,40 @@ parser.add_argument('--mnist_classes', ###################################################################### +if torch.cuda.is_available(): + device = torch.device('cuda') +else: + device = torch.device('cpu') + +###################################################################### + +def entropy(target): + probas = [] + for k in range(target.max() + 1): + n = (target == k).sum().item() + if n > 0: probas.append(n) + probas = torch.tensor(probas).float() + probas /= probas.sum() + return - (probas * probas.log()).sum().item() + +###################################################################### + args = parser.parse_args() if args.seed >= 0: torch.manual_seed(args.seed) -used_MNIST_classes = torch.tensor(eval('[' + args.mnist_classes + ']')) +used_MNIST_classes = torch.tensor(eval('[' + args.mnist_classes + ']'), device = device) ###################################################################### train_set = torchvision.datasets.MNIST('./data/mnist/', train = True, download = True) -train_input = train_set.train_data.view(-1, 1, 28, 28).float() -train_target = train_set.train_labels +train_input = train_set.train_data.view(-1, 1, 28, 28).to(device).float() +train_target = train_set.train_labels.to(device) test_set = torchvision.datasets.MNIST('./data/mnist/', train = False, download = True) -test_input = test_set.test_data.view(-1, 1, 28, 28).float() -test_target = test_set.test_labels - -if torch.cuda.is_available(): - used_MNIST_classes = used_MNIST_classes.cuda() - train_input, train_target = train_input.cuda(), train_target.cuda() - test_input, test_target = test_input.cuda(), test_target.cuda() +test_input = test_set.test_data.view(-1, 1, 28, 28).to(device).float() +test_target = test_set.test_labels.to(device) mu, std = train_input.mean(), train_input.std() train_input.sub_(mu).div_(std) @@ -58,7 +70,7 @@ test_input.sub_(mu).div_(std) # Returns a triplet of tensors (a, b, c), where a and b contain each # half of the samples, with a[i] and b[i] of same class for any i, and -# c is a 1d long tensor with the count of pairs per class used. +# c is a 1d long tensor real classes def create_image_pairs(train = False): ua, ub = [], [] @@ -76,18 +88,25 @@ def create_image_pairs(train = False): hs = x.size(0)//2 ua.append(x.narrow(0, 0, hs)) ub.append(x.narrow(0, hs, hs)) + uc.append(target[used_indices]) a = torch.cat(ua, 0) b = torch.cat(ub, 0) + c = torch.cat(uc, 0) perm = torch.randperm(a.size(0)) a = a[perm].contiguous() b = b[perm].contiguous() - c = torch.tensor([x.size(0) for x in ua]) return a, b, c ###################################################################### +# Returns a triplet a, b, c where a are the standard MNIST images, c +# the classes, and b is a Nx2 tensor, eith for every n: +# +# b[n, 0] ~ Uniform(0, 10) +# b[n, 1] ~ b[n, 0] + Uniform(0, 0.5) + c[n] + def create_image_values_pairs(train = False): ua, ub = [], [] @@ -105,21 +124,20 @@ def create_image_values_pairs(train = False): target = target[used_indices].contiguous() a = input + c = target b = a.new(a.size(0), 2) b[:, 0].uniform_(10) b[:, 1].uniform_(0.5) b[:, 1] += b[:, 0] + target.float() - c = torch.tensor([(target == k).sum().item() for k in used_MNIST_classes]) - return a, b, c ###################################################################### -class NetImagePair(nn.Module): +class NetForImagePair(nn.Module): def __init__(self): - super(NetImagePair, self).__init__() + super(NetForImagePair, self).__init__() self.features_a = nn.Sequential( nn.Conv2d(1, 16, kernel_size = 5), nn.MaxPool2d(3), nn.ReLU(), @@ -148,9 +166,9 @@ class NetImagePair(nn.Module): ###################################################################### -class NetImageValuesPair(nn.Module): +class NetForImageValuesPair(nn.Module): def __init__(self): - super(NetImageValuesPair, self).__init__() + super(NetForImageValuesPair, self).__init__() self.features_a = nn.Sequential( nn.Conv2d(1, 16, kernel_size = 5), nn.MaxPool2d(3), nn.ReLU(), @@ -180,10 +198,10 @@ class NetImageValuesPair(nn.Module): if args.data == 'image_pair': create_pairs = create_image_pairs - model = NetImagePair() + model = NetForImagePair() elif args.data == 'image_values_pair': create_pairs = create_image_values_pairs - model = NetImageValuesPair() + model = NetForImageValuesPair() else: raise Exception('Unknown data ' + args.data) @@ -195,17 +213,11 @@ print('nb_parameters %d' % sum(x.numel() for x in model.parameters())) optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3) -if torch.cuda.is_available(): - model.cuda() +model.to(device) for e in range(nb_epochs): - input_a, input_b, count = create_pairs(train = True) - - # The information bound is the entropy of the class distribution - class_proba = count.float() - class_proba /= class_proba.sum() - class_entropy = - (class_proba.log() * class_proba).sum().item() + input_a, input_b, classes = create_pairs(train = True) input_br = input_b[torch.randperm(input_b.size(0))] @@ -223,19 +235,15 @@ for e in range(nb_epochs): acc_mi /= (input_a.size(0) // batch_size) - print('%d %.04f %.04f' % (e, acc_mi / math.log(2), class_entropy / math.log(2))) + print('%d %.04f %.04f' % (e, acc_mi / math.log(2), entropy(classes) / math.log(2))) sys.stdout.flush() ###################################################################### -input_a, input_b, count = create_pairs(train = False) +input_a, input_b, classes = create_pairs(train = False) for e in range(nb_epochs): - class_proba = count.float() - class_proba /= class_proba.sum() - class_entropy = - (class_proba.log() * class_proba).sum().item() - input_br = input_b[torch.randperm(input_b.size(0))] acc_mi = 0.0 @@ -248,6 +256,6 @@ for e in range(nb_epochs): acc_mi /= (input_a.size(0) // batch_size) -print('test %.04f %.04f'%(acc_mi / math.log(2), class_entropy / math.log(2))) +print('test %.04f %.04f'%(acc_mi / math.log(2), entropy(classes) / math.log(2))) ###################################################################### -- 2.20.1