3 # @XREMOTE_HOST: elk.fleuret.org
4 # @XREMOTE_EXEC: ~/conda/bin/python
5 # @XREMOTE_PRE: ln -s ~/data/pytorch ./data
6 # @XREMOTE_PRE: killall -q -9 python || true
8 import math, sys, torch, torchvision
11 from torch.nn import functional as F
13 ######################################################################
15 # Returns a pair of tensors (a, b, c), where a and b are tensors
16 # containing each half of the samples, with a[i] and b[i] of same
17 # class for any i, and c is a 1d long tensor with the count of pairs
20 def create_pair_set(used_classes, input, target):
23 for i in used_classes:
24 used_indices = torch.arange(input.size(0), device = target.device)\
25 .masked_select(target == i.item())
26 x = input[used_indices]
27 x = x[torch.randperm(x.size(0))]
28 ua.append(x.narrow(0, 0, x.size(0)//2))
29 ub.append(x.narrow(0, x.size(0)//2, x.size(0)//2))
33 perm = torch.randperm(a.size(0))
34 a = a[perm].contiguous()
35 b = b[perm].contiguous()
36 c = torch.tensor([x.size(0) for x in ua])
40 ######################################################################
44 super(Net, self).__init__()
45 self.conv1 = nn.Conv2d(2, 32, kernel_size = 5)
46 self.conv2 = nn.Conv2d(32, 64, kernel_size = 5)
47 self.fc1 = nn.Linear(256, 200)
48 self.fc2 = nn.Linear(200, 1)
50 def forward(self, a, b):
51 x = torch.cat((a, b), 1)
52 x = F.relu(F.max_pool2d(self.conv1(x), kernel_size = 3))
53 x = F.relu(F.max_pool2d(self.conv2(x), kernel_size = 2))
54 x = x.view(x.size(0), -1)
55 x = F.relu(self.fc1(x))
59 ######################################################################
61 train_set = torchvision.datasets.MNIST('./data/mnist/', train = True, download = True)
62 train_input = train_set.train_data.view(-1, 1, 28, 28).float()
63 train_target = train_set.train_labels
65 test_set = torchvision.datasets.MNIST('./data/mnist/', train = False, download = True)
66 test_input = test_set.test_data.view(-1, 1, 28, 28).float()
67 test_target = test_set.test_labels
69 mu, std = train_input.mean(), train_input.std()
70 train_input.sub_(mu).div_(std)
71 test_input.sub_(mu).div_(std)
73 ######################################################################
75 # The information bound is the log of the number of classes in there
77 # used_classes = torch.tensor([ 0, 1, 3, 5, 6, 7, 8, 9])
78 used_classes = torch.tensor([ 3, 4, 7, 0 ])
80 nb_epochs, batch_size = 50, 100
83 optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
85 if torch.cuda.is_available():
87 train_input, train_target = train_input.cuda(), train_target.cuda()
88 test_input, test_target = test_input.cuda(), test_target.cuda()
90 for e in range(nb_epochs):
92 input_a, input_b, count = create_pair_set(used_classes, train_input, train_target)
94 class_proba = count.float()
95 class_proba /= class_proba.sum()
96 class_entropy = - (class_proba.log() * class_proba).sum().item()
98 input_br = input_b[torch.randperm(input_b.size(0))]
102 for batch_a, batch_b, batch_br in zip(input_a.split(batch_size),
103 input_b.split(batch_size),
104 input_br.split(batch_size)):
105 mi = model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log()
108 optimizer.zero_grad()
112 acc_mi /= (input_a.size(0) // batch_size)
114 print('%d %.04f %.04f'%(e, acc_mi / math.log(2), class_entropy / math.log(2)))
118 ######################################################################
120 input_a, input_b, count = create_pair_set(used_classes, test_input, test_target)
122 for e in range(nb_epochs):
123 class_proba = count.float()
124 class_proba /= class_proba.sum()
125 class_entropy = - (class_proba.log() * class_proba).sum().item()
127 input_br = input_b[torch.randperm(input_b.size(0))]
131 for batch_a, batch_b, batch_br in zip(input_a.split(batch_size),
132 input_b.split(batch_size),
133 input_br.split(batch_size)):
134 loss = - (model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log())
135 acc_mi -= loss.item()
137 acc_mi /= (input_a.size(0) // batch_size)
139 print('test %.04f %.04f'%(acc_mi / math.log(2), class_entropy / math.log(2)))
141 ######################################################################