#!/usr/bin/env python2
# -*- coding: UTF-8 -*-
# File: params_logger.py
# Date: Sun Aug 10 17:14:20 2014 -0700
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>

import numpy as np
import matplotlib.pylab as plt
import matplotlib.cm as cm
import cPickle as pickle
import gzip
import scipy.io as sio
import os

from logistic_sgd import LogisticRegression
from mlp import HiddenLayer
from convolutional_mlp import LeNetConvPoolLayer
from fixed_length_softmax import FixedLengthSoftmax
from sequence_softmax import SequenceSoftmax

name_dict = {LeNetConvPoolLayer: 'convpool',
             HiddenLayer: 'hidden',
             LogisticRegression: 'lr',
             FixedLengthSoftmax: 'fl-sm',
             SequenceSoftmax: 'ssm'}

class ParamsLogger(object):

    def __init__(self, input_shape, logdir='logs'):
        self.input_shape = input_shape
        self.logdir = logdir
        try:
            os.mkdir(self.logdir)
        except:
            pass

    def save_params(self, epoch, layers, layer_config):
        fname = os.path.join(self.logdir, "param{0:03d}.pkl.gz".format(epoch))
        res = {}

        layer_params = [x.params for x in layers]
        cnt = 0
        for layer, param, config in zip(layers, layer_params, layer_config):
            dic = {'type': name_dict[type(layer)] }
            cnt += 1
            dic.update(layer.get_params())

            res['layer' + str(cnt)] = dic
        res['input_shape'] = self.input_shape

        #sio.savemat(fname, res)
        with gzip.open(fname, 'wb') as f:
            pickle.dump(res, f, -1)

def plot_filters(obj):
    for layer_params in obj:
        if layer_params['type'] == name_dict[LeNetConvPoolLayer]:
            W = np.asarray(layer_params['W'])
            print W.shape
            filter_cnt = W.shape[0]
            n_img = W.shape[1]
            filter_size = W.shape[2]
            W = W.reshape((filter_cnt * n_img, filter_size, filter_size))
            tot_filter = filter_cnt * n_img

            # too many filters won't look good in the graph
            if len(W) > 100:
                W = W[:100]
                tot_filter = 100

            plt_num = np.ceil(np.sqrt(tot_filter))
            fig = plt.figure()
            for (idx, filt) in enumerate(W):
                if idx % 50 == 0:
                    print "Progress:", idx
                ax = fig.add_subplot(plt_num, plt_num, idx)
                plt.imshow(filt, interpolation='nearest', cmap=cm.Greys_r)
            plt.show()

def plot_filters_in_epoch(epoch, logdir='logs'):
    obj = json.loads(open(os.path.join(logdir, str(epoch) + '.json')).read())
    plot_filters(obj)

def plot_filters_with_file(filename):
    obj = json.loads(open(filename).read())
    plot_filters(obj)

if __name__ == '__main__':
    plot_filters_in_epoch(39)

@@ -91,13 +91,13 @@
  17:
-14:20
+43:11
  201
@@ -991,32 +991,122 @@
  layer_config):%0A
+        if not isinstance(epoch, basestring):%0A            epoch = %22%7B0:03d%7D%22.format(epoch)%0A
         fname = 
@@ -1139,20 +1139,16 @@
 %22param%7B0
-:03d
 %7D.pkl.gz

fix save best param