#!/usr/bin/env python
# coding: utf-8

# # Lesson 3: 'Dogs vs Cats' using Keras
# fast.ai is a machine learning library built on top of PyTorch. Keras is also a machine learning framework, but it can sit on top of several different frameworks.
# 
# In this example it sits on top of tensorflow.

# In[1]:


get_ipython().run_line_magic('reload_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


# First we set up our data paths, like normal
PATH = "data/dogscats/"
sz=224
batch_size=64


# In[3]:


# Then we import keras
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing import image
from keras.layers import Dropout, Flatten, Dense
from keras.applications import ResNet50
from keras.models import Model, Sequential
from keras.layers import Dense, GlobalAveragePooling2D
from keras import backend as K
from keras.applications.resnet50 import preprocess_input


# In[4]:


# We must set up paths to our training data and our validation data
train_data_dir = f'{PATH}train'
validation_data_dir = f'{PATH}valid'


# In[5]:


# Following that, we have to set up "data generators" for both our training and test (validation) data.
# These specify the transformations and normalization required

train_datagen = ImageDataGenerator(preprocessing_function=preprocess_input,
    shear_range=0.2, zoom_range=0.2, horizontal_flip=True)

test_datagen = ImageDataGenerator(preprocessing_function=preprocess_input)

# We can now declare our data generators.
# class_mode='binary' means we are going to classify between 2 different classes.
# Otherwise we could use categorical
train_generator = train_datagen.flow_from_directory(train_data_dir,
    target_size=(sz, sz),
    batch_size=batch_size, class_mode='binary')

# Note here the shuffle parameter. For training, we shuffle the input to make things as random as possible
# But for validation, we want to keep the order of the images so we don't mess up the labels
validation_generator = test_datagen.flow_from_directory(validation_data_dir,
    shuffle=False,
    target_size=(sz, sz),
    batch_size=batch_size, class_mode='binary')


# In[6]:


# Now we create a base model and declare the different layers we want to use
base_model = ResNet50(weights='imagenet', include_top=False)
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x)
predictions = Dense(1, activation='sigmoid')(x)


# In[7]:


# Having set up our base model, we set up our full model
model = Model(inputs=base_model.input, outputs=predictions)

# We have to manually freeze the layers one-by-one
for layer in base_model.layers: layer.trainable = False

# Models must be compiled before they can be used    
model.compile(optimizer='rmsprop', loss='binary_crossentropy', metrics=['accuracy'])


# In[8]:


# Do the training
# (Interrupted due to time)
get_ipython().run_line_magic('%time', '')
model.fit_generator(train_generator, train_generator.n // batch_size, epochs=3, workers=4,
        validation_data=validation_generator, validation_steps=validation_generator.n // batch_size)


# In[9]:


# There is no support for differential learning rates
# So we split the model at 140, freeze all before, and unfreeze all layers after
# Then train again
split_at = 140
for layer in model.layers[:split_at]: layer.trainable = False
for layer in model.layers[split_at:]: layer.trainable = True
model.compile(optimizer='rmsprop', loss='binary_crossentropy', metrics=['accuracy'])    


# In[10]:


get_ipython().run_cell_magic('time', '', 'model.fit_generator(train_generator, train_generator.n // batch_size, epochs=1, workers=3,\n        validation_data=validation_generator, validation_steps=validation_generator.n // batch_size)\n')


# In[ ]: