Update.

diff --git a/mine_mnist.py b/mine_mnist.py
index f573b89..c6dc287 100755 (executable)
--- a/mine_mnist.py
+++ b/mine_mnist.py
@@ -3,7 +3,7 @@
 # @XREMOTE_HOST: elk.fleuret.org
 # @XREMOTE_EXEC: ~/conda/bin/python
 # @XREMOTE_PRE: ln -s ~/data/pytorch ./data
-# @XREMOTE_PRE: killall -q -9 python || echo "Nothing killed"
+# @XREMOTE_PRE: killall -q -9 python || true
 
 import math, sys, torch, torchvision
 
@@ -12,9 +12,9 @@ from torch.nn import functional as F
 
 ######################################################################
 
-# Returns a pair of tensors (x, c), where x is a Nx2x28x28 containing
-# pairs of images of same classes (one per channel), and p is a 1d
-# long tensor with the count of pairs per class used
+# Returns a pair of tensors (a, b, c), where a and b are Nx1x28x28
+# tensors containing images, 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.
 
 def create_pair_set(used_classes, input, target):
     u = []
@@ -28,10 +28,11 @@ def create_pair_set(used_classes, input, target):
         x = x[0:2 * (x.size(0) // 2)].reshape(-1, 2, 28, 28)
         u.append(x)
 
-    x = torch.cat(u, 0).contiguous()
+    x = torch.cat(u, 0)
+    x = x[torch.randperm(x.size(0))]
     c = torch.tensor([x.size(0) for x in u])
 
-    return x, c
+    return x.narrow(1, 0, 1).contiguous(), x.narrow(1, 1, 1).contiguous(), c
 
 ######################################################################
 
@@ -43,7 +44,8 @@ class Net(nn.Module):
         self.fc1 = nn.Linear(256, 200)
         self.fc2 = nn.Linear(200, 1)
 
-    def forward(self, x):
+    def forward(self, a, b):
+        x = torch.cat((a, b), 1)
         x = F.relu(F.max_pool2d(self.conv1(x), kernel_size = 3))
         x = F.relu(F.max_pool2d(self.conv2(x), kernel_size = 2))
         x = x.view(x.size(0), -1)
@@ -57,8 +59,13 @@ train_set = torchvision.datasets.MNIST('./data/mnist/', train = True, download =
 train_input  = train_set.train_data.view(-1, 1, 28, 28).float()
 train_target = train_set.train_labels
 
+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
+
 mu, std = train_input.mean(), train_input.std()
 train_input.sub_(mu).div_(std)
+test_input.sub_(mu).div_(std)
 
 ######################################################################
 
@@ -75,31 +82,56 @@ optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
 if torch.cuda.is_available():
     model.cuda()
     train_input, train_target = train_input.cuda(), train_target.cuda()
+    test_input, test_target = test_input.cuda(), test_target.cuda()
 
 for e in range(nb_epochs):
-    input, count = create_pair_set(used_classes, train_input, train_target)
+
+    input_a, input_b, count = create_pair_set(used_classes, train_input, train_target)
 
     class_proba = count.float()
     class_proba /= class_proba.sum()
     class_entropy = - (class_proba.log() * class_proba).sum().item()
 
-    input = input[torch.randperm(input.size(0))]
-    indep_input = input.clone()
-    indep_input[:, 1] = input[torch.randperm(input.size(0)), 1]
+    input_br = input_b[torch.randperm(input_b.size(0))]
 
     mi = 0.0
 
-    for batch, indep_batch in zip(input.split(batch_size), indep_input.split(batch_size)):
-        loss = - (model(batch).mean() - model(indep_batch).exp().mean().log())
+    for batch_a, batch_b, batch_br in zip(input_a.split(batch_size),
+                                          input_b.split(batch_size),
+                                          input_br.split(batch_size)):
+        loss = - (model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log())
         mi -= loss.item()
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
 
-    mi /= (input.size(0) // batch_size)
+    mi /= (input_a.size(0) // batch_size)
 
     print('%d %.04f %.04f'%(e, mi / math.log(2), class_entropy / math.log(2)))
 
     sys.stdout.flush()
 
 ######################################################################
+
+input_a, input_b, count = create_pair_set(used_classes, test_input, test_target)
+
+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))]
+
+    mi = 0.0
+
+    for batch_a, batch_b, batch_br in zip(input_a.split(batch_size),
+                                          input_b.split(batch_size),
+                                          input_br.split(batch_size)):
+        loss = - (model(batch_a, batch_b).mean() - model(batch_a, batch_br).exp().mean().log())
+        mi -= loss.item()
+
+    mi /= (input_a.size(0) // batch_size)
+
+print('test %.04f %.04f'%(mi / math.log(2), class_entropy / math.log(2)))
+
+######################################################################