from theano.sandbox import cuda

%matplotlib inline
import utils; reload(utils)
from utils import *
from __future__ import division, print_function

#path = os.path.join('input','sample')
path = os.path.join('input','sample-10')
#path = os.path.join('input')
output_path = os.path.join('output','sample')
model_path = os.path.join(output_path, 'models')
if not os.path.exists(model_path): os.mkdir(model_path)

#batch_size=64
#batch_size=32
#batch_size=16
batch_size=8

"""
Explanation: Training a better model
End of explanation
"""


model = vgg_ft(2)

??vgg_ft

"""
Explanation: Are we underfitting?
Our validation accuracy so far has generally been higher than our training accuracy. That leads to two obvious questions:

How is this possible?
Is this desirable?

The answer to (1) is that this is happening because of dropout. Dropout refers to a layer that randomly deletes (i.e. sets to zero) each activation in the previous layer with probability p (generally 0.5). This only happens during training, not when calculating the accuracy on the validation set, which is why the validation set can show higher accuracy than the training set.
The purpose of dropout is to avoid overfitting. By deleting parts of the neural network at random during training, it ensures that no one part of the network can overfit to one part of the training set. The creation of dropout was one of the key developments in deep learning, and has allowed us to create rich models without overfitting. However, it can also result in underfitting if overused, and this is something we should be careful of with our model.
So the answer to (2) is: this is probably not desirable. It is likely that we can get better validation set results with less (or no) dropout, if we're seeing that validation accuracy is higher than training accuracy - a strong sign of underfitting. So let's try removing dropout entirely, and see what happens!
(We had dropout in this model already because the VGG authors found it necessary for the imagenet competition. But that doesn't mean it's necessary for dogs v cats, so we will do our own analysis of regularization approaches from scratch.)
Removing dropout
Our high level approach here will be to start with our fine-tuned cats vs dogs model (with dropout), then fine-tune all the dense layers, after removing dropout from them. The steps we will take are:
- Re-create and load our modified VGG model with binary dependent (i.e. dogs v cats)
- Split the model between the convolutional (conv) layers and the dense layers
- Pre-calculate the output of the conv layers, so that we don't have to redundently re-calculate them on every epoch
- Create a new model with just the dense layers, and dropout p set to zero
- Train this new model using the output of the conv layers as training data.
Start with Vgg + binary output (dogs/cats)
As before we need to start with a working model, so let's bring in our working VGG 16 model and change it to predict our binary dependent...
End of explanation
"""


#model.load_weights(model_path+'finetune3.h5')
model.load_weights(os.path.join(model_path, 'finetune_1_ll.h5'))

"""
Explanation: ...and load our fine-tuned weights.
End of explanation
"""


layers = model.layers

last_conv_idx = [index for index,layer in enumerate(layers) 
                     if type(layer) is Convolution2D][-1]

last_conv_idx

layers[last_conv_idx]

conv_layers = layers[:last_conv_idx+1] # convolutional layers only; i.e. first N layers to the index
conv_model = Sequential(conv_layers)
# Dense layers - also known as fully connected or 'FC' layers
fc_layers = layers[last_conv_idx+1:] # remaining layers are Dense/FC

"""
Explanation: Split conv and dense layers
We're going to be training a number of iterations without dropout, so it would be best for us to pre-calculate the input to the fully connected layers - i.e. the Flatten() layer. We'll start by finding this layer in our model, and creating a new model that contains just the layers up to and including this layer:
End of explanation
"""


batches = get_batches(os.path.join(path,'train'), shuffle=False, batch_size=batch_size)
val_batches = get_batches(os.path.join(path,'valid'), shuffle=False, batch_size=batch_size)

val_classes = val_batches.classes
trn_classes = batches.classes
val_labels = onehot(val_classes)
trn_labels = onehot(trn_classes)

"""
Explanation: Generate features for the FC layers by precalculating conv output
Now we can use the exact same approach to creating features as we used when we created the linear model from the imagenet predictions in the last lesson - it's only the model that has changed. As you're seeing, there's a fairly small number of "recipes" that can get us a long way!
End of explanation
"""


val_features = conv_model.predict_generator(val_batches, val_batches.nb_sample)

trn_features = conv_model.predict_generator(batches, batches.nb_sample)

save_array(os.path.join(model_path, 'train_convlayer_features.bc'), trn_features)
save_array(os.path.join(model_path,'valid_convlayer_features.bc'), val_features)

??save_array

trn_features = load_array(os.path.join(model_path, 'train_convlayer_features.bc'))
val_features = load_array(os.path.join(model_path,'valid_convlayer_features.bc'))

trn_features.shape

val_features.shape

"""
Explanation: Below:
We're pre-calculating the inputs to the new model. The inputs are the training and valiation sets. So, we basically want to get the result of running those two sets thru the conv layers only and save them so we can run them through the new model.
So, we use the conv-only model to run prediction on the training and validation sets in order to precompute the values we need.
End of explanation
"""


# Copy the weights from the pre-trained model.
# NB: Since we're removing dropout, we want to half the weights
def proc_wgts(layer): return [o/2 for o in layer.get_weights()]

# Such a finely tuned model needs to be updated very slowly!
opt = RMSprop(lr=0.00001, rho=0.7)

def get_fc_model():
    model = Sequential([
        MaxPooling2D(input_shape=conv_layers[-1].output_shape[1:]),
        Flatten(),
        Dense(4096, activation='relu'),
        Dropout(0.),
        Dense(4096, activation='relu'),
        Dropout(0.),
        Dense(2, activation='softmax')
        ])

    for l1,l2 in zip(model.layers, fc_layers): l1.set_weights(proc_wgts(l2))

    model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
    return model

fc_model = get_fc_model()

"""
Explanation: Remove dropout from the fully-connected layer model
For our new fully connected model, we'll create it using the exact same architecture as the last layers of VGG 16, so that we can conveniently copy pre-trained weights over from that model. However, we'll set the dropout layer's p values to zero, so as to effectively remove dropout.
End of explanation
"""


fc_model.fit(trn_features, trn_labels, nb_epoch=8, 
             batch_size=batch_size, validation_data=(val_features, val_labels))

??fc_model.fit

"""
Explanation: Fit the FC model to the training and validation data
And fit the model in the usual way:
End of explanation
"""


fc_model.save_weights(os.path.join(model_path,'lesson3_no_dropout.h5'))

fc_model.load_weights(os.path.join(model_path,'lesson3_no_dropout.h5'))

"""
Explanation: Save the weights (no dropout)
End of explanation
"""


gen = image.ImageDataGenerator(rotation_range=15, width_shift_range=0.1, 
                               height_shift_range=0.1, zoom_range=0.1, horizontal_flip=True)

batches = get_batches(os.path.join(path, 'train'), gen, batch_size=batch_size)
# NB: We don't want to augment or shuffle the validation set
val_batches = get_batches(os.path.join(path, 'valid'), shuffle=False, batch_size=batch_size)

"""
Explanation: Reducing overfitting
Now that we've gotten the model to overfit, we can take a number of steps to reduce this.
Add data augmentation to the training data
Let's try adding a small amount of data augmentation, and see if we reduce overfitting as a result. The approach will be identical to the method we used to finetune the dense layers in lesson 2, except that we will use a generator with augmentation configured. Here's how we set up the generator, and create batches from it:
End of explanation
"""


fc_model = get_fc_model()

for layer in conv_model.layers: layer.trainable = False
# Look how easy it is to connect two models together!
conv_model.add(fc_model)

"""
Explanation: Combine the Conv and FC layers into a single model
When using data augmentation, we can't pre-compute our convolutional layer features, since randomized changes are being made to every input image. That is, even if the training process sees the same image multiple times, each time it will have undergone different data augmentation, so the results of the convolutional layers will be different.
Therefore, in order to allow data to flow through all the conv layers and our new dense layers, we attach our fully connected model to the convolutional model--after ensuring that the convolutional layers are not trainable:
End of explanation
"""


conv_model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])

conv_model.fit_generator(batches, samples_per_epoch=batches.nb_sample, nb_epoch=8, 
                        validation_data=val_batches, nb_val_samples=val_batches.nb_sample)

"""
Explanation: Now we can compile, train, and save our model as usual - note that we use fit_generator() since we want to pull random images from the directories on every batch.
Compile and train the combined model on augmented data
End of explanation
"""


conv_model.save_weights(os.path.join(model_path, 'aug1.h5'))

conv_model.load_weights(os.path.join(model_path, 'aug1.h5'))

"""
Explanation: Save the weights (combined)
End of explanation
"""


conv_layers[-1].output_shape[1:]  # last layer shape

def get_bn_layers(p):
    return [
        MaxPooling2D(input_shape=conv_layers[-1].output_shape[1:]),
        Flatten(),
        Dense(4096, activation='relu'),
        BatchNormalization(),
        Dropout(p),
        Dense(4096, activation='relu'),
        BatchNormalization(),
        Dropout(p),
        Dense(1000, activation='softmax')
        ]

def load_fc_weights_from_vgg16bn(model):
    "Load weights for model from the dense layers of the Vgg16BN model."
    # See imagenet_batchnorm.ipynb for info on how the weights for
    # Vgg16BN can be generated from the standard Vgg16 weights.
    from vgg16bn import Vgg16BN
    vgg16_bn = Vgg16BN()
    _, fc_layers = split_at(vgg16_bn.model, Convolution2D)
    copy_weights(fc_layers, model.layers)

p=0.6

"""
Explanation: Add batchnorm to the model
We can use nearly the same approach as before - but this time we'll add batchnorm layers (and dropout layers):
End of explanation
"""


#bn_model = Sequential(get_bn_layers(0.6))
bn_model = Sequential(get_bn_layers(p))

load_fc_weights_from_vgg16bn(bn_model)

def proc_wgts(layer, prev_p, new_p):
    scal = (1-prev_p)/(1-new_p)
    return [o*scal for o in layer.get_weights()]

for l in bn_model.layers: 
    if type(l)==Dense: l.set_weights(proc_wgts(l, 0.5, 0.6))  # here l == 'l' not '1'

# Remove last layer and lock all the others
bn_model.pop()
for layer in bn_model.layers: layer.trainable=False

# Add linear layer (2-class) (just doing the ImageNet mapping to Kaggle dogs and cats)
bn_model.add(Dense(2,activation='softmax'))

"""
Explanation: Create a standalone model from the BN layers of vgg16bn
End of explanation
"""


bn_model.compile(Adam(), 'categorical_crossentropy', metrics=['accuracy'])

bn_model.fit(trn_features, trn_labels, nb_epoch=8, validation_data=(val_features, val_labels))

"""
Explanation: Compile and fit the model
End of explanation
"""


bn_model.save_weights(os.path.join(model_path,'bn.h5'))

bn_model.load_weights(os.path.join(model_path,'bn.h5'))

"""
Explanation: Save the weights (batchnorm)
End of explanation
"""


bn_layers = get_bn_layers(0.6)
bn_layers.pop()
bn_layers.append(Dense(2,activation='softmax'))

final_model = Sequential(conv_layers)
for layer in final_model.layers: layer.trainable = False
for layer in bn_layers: final_model.add(layer)

"""
Explanation: Create another BN model and combine it with the conv layers into a final model
End of explanation
"""


for l1,l2 in zip(bn_model.layers, bn_layers):
    l2.set_weights(l1.get_weights())

"""
Explanation: Set the BN layers weights from the first BN model
End of explanation
"""


final_model.compile(optimizer=Adam(), 
                    loss='categorical_crossentropy', metrics=['accuracy'])

final_model.fit_generator(batches, samples_per_epoch=batches.nb_sample, nb_epoch=1, 
                        validation_data=val_batches, nb_val_samples=val_batches.nb_sample)

"""
Explanation: Fit the model
End of explanation
"""


final_model.save_weights(os.path.join(model_path, 'final1.h5'))

"""
Explanation: Save the weights (final model)
End of explanation
"""


final_model.fit_generator(batches, samples_per_epoch=batches.nb_sample, nb_epoch=4, 
                        validation_data=val_batches, nb_val_samples=val_batches.nb_sample)

"""
Explanation: Fit the model
End of explanation
"""


final_model.save_weights(os.path.join(model_path, 'final2.h5'))

final_model.optimizer.lr=0.001

"""
Explanation: Save the weights (final model)
End of explanation
"""


final_model.fit_generator(batches, samples_per_epoch=batches.nb_sample, nb_epoch=4, 
                        validation_data=val_batches, nb_val_samples=val_batches.nb_sample)

"""
Explanation: Fit the model
End of explanation
"""


bn_model.save_weights(os.path.join(model_path, 'final3.h5'))

"""
Explanation: Save the weights (final model)
End of explanation
"""


