IIC MNIST inference

mnist_inf.py

from code.utils.cluster.cluster_eval import cluster_eval
from code.utils.cluster.data import cluster_twohead_create_dataloaders
import pickle
import code.archs as archs
import torch
from datetime import datetime
import sys

config_in = open("/scratch/shared/nfs1/xuji/iid_private/685/config.pickle", "rb")
config = pickle.load(config_in)
net = archs.__dict__[config.arch](config)
net.load_state_dict(torch.load("/scratch/shared/nfs1/xuji/iid_private/685/best_net.pytorch"))
net.cuda()

#config.batch_sz = 100

dataloaders_head_A, dataloaders_head_B, \
mapping_assignment_dataloader, mapping_test_dataloader = \
  cluster_twohead_create_dataloaders(config)

if "MNIST" in config.dataset:
  sobel = False
else:
  sobel = True


"""
cluster_eval(config, net,
             mapping_assignment_dataloader=mapping_assignment_dataloader,
             mapping_test_dataloader=mapping_test_dataloader,
             sobel=sobel, print_stats=True)
"""

best_head = 0
mappings = dict([(0, 9), (1, 3), (2, 1), (3, 4), (4, 7), (5, 8), (6, 5), (7, 6), (8, 0), (9, 2)])

acc = 0.
ct = 0

net.eval()

for i, batch in enumerate(mapping_test_dataloader):
  print("batch %d %s" % (i, datetime.now()))
  sys.stdout.flush()

  imgs = batch[0].cuda()
  labels = batch[1].cuda()

  preds = net(imgs)[best_head]
  preds_flat = preds.argmax(dim=1, keepdim=False)
  preds_flat_reordered = torch.zeros(preds_flat.shape, dtype=preds_flat.dtype).cuda()
  for pred_c, target_c in mappings.iteritems():
    samples_pred_c = preds_flat == pred_c
    preds_flat_reordered[samples_pred_c] = target_c

  acc += (preds_flat_reordered == labels).sum().item()
  ct += labels.shape[0]

print((acc, ct)) # prints (69477.0, 70000)
print(acc/ float(ct)) # prints 0.992528571429