<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Kaggle Dogs Vs. Cats Using LeNet  on Google Colab TPU
Required setup
Update api_token with kaggle api key for downloading dataset
  - Login to kaggle
  - My Profile &gt; Edit Profile &gt; Createt new API Token
  - Update **api_token** dict below with the values
Change Notebook runtime to TPU
  - In colab notebook menu, Runtime &gt; Change runtime type
  - Select TPU in the list
Install kaggle package, download and extract zip file
Step1:  Re-arrange classes to 2 separate directories
Step2:  Training configs
Step3:  Setup generators to provide with train and validation batches
Step4:  Define LeNet model architecture
Step5:  Check for TPU availability
Step6:  Convert keras model to TPU model
Step7:  Run training
Step8:  Save the model weights
Step9:  Download model weights locally
<ASSISTANT_TASK:>
Python Code:
!pip install kaggle
api_token = {"username":"xxxxx","key":"xxxxxxxxxxxxxxxxxxxxxxxx"}
import json
import zipfile
import os
os.mkdir('/root/.kaggle')
with open('/root/.kaggle/kaggle.json', 'w') as file:
    json.dump(api_token, file)
!chmod 600 /root/.kaggle/kaggle.json
# !kaggle config path -p /root
!kaggle competitions download -c dogs-vs-cats
zip_ref = zipfile.ZipFile('/content/train.zip', 'r')
zip_ref.extractall()
zip_ref.close()
Explanation: Kaggle Dogs Vs. Cats Using LeNet  on Google Colab TPU
Required setup
Update api_token with kaggle api key for downloading dataset
  - Login to kaggle
  - My Profile &gt; Edit Profile &gt; Createt new API Token
  - Update **api_token** dict below with the values
Change Notebook runtime to TPU
  - In colab notebook menu, Runtime &gt; Change runtime type
  - Select TPU in the list
Install kaggle package, download and extract zip file
End of explanation
!mkdir train/cat train/dog
!mv train/*cat*.jpg train/cat
!mv train/*dog*.jpg train/dog
Explanation: Re-arrange classes to 2 separate directories
End of explanation
BATCH_SIZE   = 64
IMG_DIM      = (256, 256, 3)
NUM_EPOCHS   = 1
Explanation: Training configs
End of explanation
import tensorflow as tf
from tensorflow import keras
print(keras.__version__)
print(tf.__version__)
datagen = keras.preprocessing.image.ImageDataGenerator(
    rescale=1./255,
    validation_split=0.2)
traingen = datagen.flow_from_directory(
    'train',
    batch_size = BATCH_SIZE,
    target_size = IMG_DIM[:-1],
    class_mode = 'categorical',
    subset='training')
valgen = datagen.flow_from_directory(
    'train',
    batch_size = BATCH_SIZE,
    target_size = IMG_DIM[:-1],
    class_mode = 'categorical',
    subset='validation')
Explanation: Setup generators to provide with train and validation batches
End of explanation
input = keras.layers.Input(IMG_DIM, name="input")
conv1 = keras.layers.Conv2D(20, kernel_size=(5, 5), padding='same')(input)
pool1 = keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2))(conv1)
conv2 = keras.layers.Conv2D(50, kernel_size=(5,5), padding='same')(pool1)
pool2 = keras.layers.MaxPooling2D(pool_size=(2,2), strides=(2,2))(conv1)
flatten1 = keras.layers.Flatten()(pool2)
fc1 = keras.layers.Dense(500, activation='relu')(flatten1)
fc2 = keras.layers.Dense(2, activation='softmax')(fc1)
model = keras.models.Model(inputs=input, outputs=fc2)
model.compile(
    loss='categorical_crossentropy',
    optimizer=keras.optimizers.SGD(lr=0.01),
    metrics=['accuracy'])
print(model.summary())
Explanation: Define LeNet model architecture
End of explanation
import os
try:
  device_name = os.environ['COLAB_TPU_ADDR']
  TPU_ADDRESS = 'grpc://' + device_name
  print('Found TPU at: {}'.format(TPU_ADDRESS))
except KeyError:
  print('TPU not found')
Explanation: Check for TPU availability
End of explanation
tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_ADDRESS)))
Explanation: Convert keras model to TPU model
End of explanation
tpu_model.fit_generator(
    traingen,
    steps_per_epoch=traingen.n//traingen.batch_size,
    epochs=1,
    validation_data=valgen,
    validation_steps=valgen.n//valgen.batch_size)
Explanation: Run training
End of explanation
tpu_model.save_weights('./lenet-catdog.h5', overwrite=True)
Explanation: Save the model weights
End of explanation
from google.colab import files
files.download("lenet-catdog.h5")
Explanation: Download model weights locally
End of explanation
<END_TASK>