h_classifier_2.py

import os, sys
os.environ['GLOG_minloglevel'] = '2'    # suprress Caffe verbose prints

import settings
sys.path.insert(0, settings.caffe_root)
import caffe

import numpy as np
from numpy.linalg import norm
import scipy.misc, scipy.io
import argparse 
import util
from sampler import Sampler

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class MLP(nn.Module):

    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(4096, 4096)
        self.fc2 = nn.Linear(4096, 1000)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x       

def get_code(generator, classifier, batch_size):
    '''
    Generate <batch_size> h's corresponding to images generated by generator. 
    Sample uniformly from h distribution, feed through generator and classifier, and if prob_highest_class > 55% then choose h.
    return h and its corresponding class
    '''
    gen_in = settings.generator_in_layer
    gen_out = settings.generator_out_layer
    h_shape = generator.blobs[gen_in].data.shape
    
    # Get the input and output sizes
    image_shape = classifier.blobs['data'].data.shape
    generator_output_shape = generator.blobs[gen_out].data.shape

    # Calculate the difference between the input image of the condition net 
    # and the output image from the generator
    image_size = util.get_image_size(image_shape)
    generator_output_size = util.get_image_size(generator_output_shape)

    # The top left offset to crop the output image to get a 227x227 image
    topleft = util.compute_topleft(image_size, generator_output_size)
    
    h_list = []
    class_list = []
    i = 0
    print ("starting to generate h's")
    while len(h_list) < batch_size:
        print ("round %d" %i)
        # Sample h from uniform distribution
        h = np.random.normal(0, 1, h_shape)
        
        # Push h through Generator to get image
        generator.blobs[gen_in].data[:] = h
        generated = generator.forward()
        x = generated[gen_out].copy()       # 256x256

        # Crop from 256x256 to 227x227
        cropped_x = x[:,:,topleft[0]:topleft[0]+image_size[0], topleft[1]:topleft[1]+image_size[1]]
        cropped_x_copy = cropped_x.copy()
        
        softmax_output = classifier.forward(data=cropped_x_copy, end='prob')
        probs = softmax_output['prob'][0]
        best_class = probs.argmax()     # highest probability unit
        print (probs[best_class])
        if probs[best_class] > 55:
            h_list.append(h)
            class_list.append(best_class)
            i +=1
    return h

if __name__ == '__main__':
    # parameters
    num_iter = 10
    batch_size = 50
    lr = 0.01
    
    classifier = caffe.Classifier(settings.encoder_definition, settings.encoder_weights,
                             mean = np.float32([104.0, 117.0, 123.0]), # ImageNet mean
                             channel_swap = (2,1,0)) # the reference model has channels in BGR order instead of RGB
    generator = caffe.Net(settings.generator_definition, settings.generator_weights, caffe.TEST)
    

    # train h_classifier
    h_classifier = MLP()
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(h_classifier.parameters(), lr=lr, betas=(0.9,0.999), eps=1e-08, weight_decay=0.0005)
    
    print("Starting Training")
    for i in range(num_iter):
        print("iter %d" %i)
        running_loss = 0.0
        
        # get the batch 
        hs, labels = get_code(generator, classifier, batch_size)
        hs, labels = Variable(hs), Variable(labels)
        
        optimizer.zero_grad()
        
        # forward + backward + optimize
        outputs = h_classifier(hs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        #print statistics
        running_loss += loss.data[0]
        if i % 100 == 0:
            print('[%d] loss: %.3f' %
                  (i + 1, running_loss / 2000))
            running_loss = 0.0
            
    print('Finished Training')