{"cells":[{"metadata":{},"cell_type":"markdown","source":"TPUs and You! \n\nIn general, you should consider TPUs if you already have a solid pipeline with a GPU and are looking to perform a hyper parameter search. TPUs will crunch your data so much faster because the hardware has cores designed for matrix calculations! No more scalar calculations. No more vector calculations. You go straight into the matrix! \n\nHere's a link to the reference docs. https://cloud.google.com/tpu/docs/concepts \n\nThis notebook will cover the tfrecord format and some starter code for it. How you build you dataset will determine how fast your TPU operates. \n\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import tensorflow as tf \nimport tensorflow.keras as tfk\nimport numpy as np \nimport glob\nimport pandas as pd \nfrom skimage import io, transform\nfrom tqdm.notebook import tqdm\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n\nmetadata = pd.read_csv('/kaggle/input/deepfake-first-frames-and-labels/train_df.csv')\nmetadata['path'] = '/kaggle/input/deepfake-first-frames-and-labels/reduced_train_stills/first_stills/' + metadata['index']\nmetadata['path'] = metadata.path.str.replace('.mp4', '.jpg')\nmetadata.head(5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Training on a TPU has a completely different set of problems than a GPU. In TPU world, the CPU is often the bottleneck. This is because each TPU node has 8 super fast matrix calculation cores that will crunch through your data at a crazy speed. You can augment your data during TPU training, but you need to be careful about which augmentations you use. Augmentations from tf.image and tfa.image generally work because they interact with tensors in a graph. If your augmentations do not interact with tensors in a graph, they will likely use the CPU of the TPU node, which can create a bottleneck. \n\nData loading has huge impact on TPU throughput. As an example, when I went from reading tfrecord datasets in serial to reading them in parallel, it doubled my TPU throughput.\n\nBelow is some starter code for building and reading TFrecords. \n"},{"metadata":{"trusted":true},"cell_type":"code","source":"#bytes feature is the easiest feature type to use. It is also extremely inefficient in terms of storage. This one is configured for images\n#you can save on storage by compressing your tensor into a jpg byte stream \n#you can also compress the whole tfrecord file, but I haven't tried that yet \n\ndef _bytes_feature(value):\n    value = tf.io.serialize_tensor(tf.convert_to_tensor(value, dtype = tf.uint8))\n    if tf.executing_eagerly: value = value.numpy()\n    else: value = value.eval() \n    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))\n\n# this method takes an image and label and converts it to a bytes feature then serializes it \n# you can customize this and make it more efficient.\ndef serialize_example(image, label):\n    \n    feature = {'image' : _bytes_feature(image)\n              ,'label' : _bytes_feature(label)\n              }\n    \n    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n    return example_proto.SerializeToString()\n\n#this feature map is very important. It allows the TPU to parse the tfrecord and will be used later \nfeature_map = {'image' :tf.io.FixedLenFeature([], tf.string, default_value='')\n              ,'label' : tf.io.FixedLenFeature([], tf.string, default_value='')\n              }\n\ndef get_xy(metadatarow):\n    x = io.imread(metadatarow.path)\n    \n    #you can transform your image here \n    x = transform.resize(x, (256,256))\n    \n    #please remember that fakes = 1 and reals = 0\n    y = 1\n    if metadatarow.label == 'REAL': y = 0\n    return x, y \n    \n#these record files are really small. I recommend using a sample number that gets you to ~150mb per tfrecord file\nsamples = 64\n\n\n#you can generate you own filenames \nfilenames = ['test1.tfrecord', 'test2.tfrecord']\n\nfor filename in filenames:     \n    with tf.io.TFRecordWriter(filename) as writer:         \n        for i, row in  tqdm(metadata.sample(n = samples).iterrows(), total = samples): \n            X, y = get_xy(row)\n            example = serialize_example(X, y)\n            writer.write(example)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n#parsing an example requires a feature map, as defined earlier\n#remember that whatever things you did to get the X, y pairs into the tfrecord, you will need to undo them so you get your xy pairs back \ndef parse_example(stringlike): \n    parsed = tf.io.parse_single_example(stringlike, feature_map)\n    return (tf.io.parse_tensor(parsed['image'], out_type = tf.uint8)\n            , tf.io.parse_tensor(parsed['label'], out_type = tf.uint8))\n\n#tfrecords are streams of data, so you will need to remind the TPU what shape the data is in \ndef force_shape(tx, ty): \n    xshape = (256, 256, 3)\n    yshape = (1, 1)\n    tx = tf.reshape(tx, xshape)\n    ty = tf.reshape(ty, yshape)\n    return tx, ty\n\n#this is a helper function that I made that allows you to wrap a mappable function, so that you can pass parameters to it \ndef get_mappable(func, **kwargs):\n    #sets kwargs for a function that should map across a dataset\n    default_kwargs = kwargs\n    def map_wrap(*args):\n        return func(*args, **default_kwargs)\n    \n    return map_wrap\n\ndef force_shape_parametered(tx, ty, xshape, yshape): \n    tx = tf.reshape(tx, xshape)\n    ty = tf.reshape(ty, yshape)\n    return tx, ty\n\n#VERY IMPORTANT. OMG THIS WAS SO HARD TO FIND. TPUs don't take your regular data types. they are very picky. For this comp, int32 is fine. \ndef force_int32(tx, ty): \n    tx = tf.dtypes.cast(tx, tf.int32)\n    ty = tf.dtypes.cast(ty, tf.int32)\n    return tx, ty \n\n#read the record. You should upload it to a cloud storage bucket\n#storage bucket paths look like gs://bucketname/something/something/test1.tfrecord \n#you can pass a lot of filenames for optimal performance \ndataset = tf.data.TFRecordDataset(filenames)\n\n#the map method applies a function across the whole dataset (technically it only applies as it reads the data)\ndataset = dataset.map(parse_example)\ndataset = dataset.map(force_int32)\n\n# a basic force shape \n# dataset = dataset.map(force_shape)\n\n# a parametered force shape, so you can define the parameters \ndataset = dataset.map(get_mappable(force_shape_parametered, xshape = [256, 256, 3], yshape = [1]))\n\n\n#you need to batch the dataset. remember that this is the batch size passed to the tpu. For this dataset, 64 per TPU core is fine. You can probably do 128. \n#remember to drop the remainder or else you will get a silly error that stops your tpu, but I believe this drop remainder is automatic now \ndataset = dataset.batch(16, drop_remainder = True)\n#prefetch also seems to be automatically tuned \ndataset = dataset.prefetch(2)\n#don't forget to repeat your datasets or else they will be exhausted\ndataset = dataset.repeat()\n\n# unfortuantely, you can only read tfrecord when it is in the cloud. below is a for loop that I use to verify the shape, before passing it on \n# you should run this loop in TPU off mode when generating data \n\n# for record in dataset.take(1): \n#     print('batch shape', record[0].shape, record[1].shape)    \n#     break","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Building the model is fairly straight forward. Whatever you were doing in tf keras before, use with strategy.scope and you will be fine. "},{"metadata":{"trusted":true},"cell_type":"code","source":"# detect and init the TPU\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)\n\n\n    \n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learning_rate = 0.01        \n\nwith strategy.scope():\n    \n    model = tfk.models.Sequential()\n    \n    pt = tfk.applications.resnet.ResNet50 \n\n    ptmod = pt(include_top=False\n                , weights='imagenet'\n                , input_tensor=None\n                , input_shape=(256, 256, 3)\n                , pooling = 'avg')\n\n\n    model.add(ptmod)\n    model.add(tfk.layers.Dropout(rate = 0.5))\n    model.add(tfk.layers.Dense(1))\n    model.add(tfk.layers.Activation('sigmoid'))\n    \n    optimizer = tfk.optimizers.Adam(learning_rate = learning_rate)\n\n\n    model.compile(loss='binary_crossentropy', optimizer=optimizer)\n\n    model.summary()\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is the training loop for a TPU. \n\nOnce you figure out how you will create your tfrecords, host them on a bucket, and then read back the tfrecords into a dataset, it is very easy to train. \n\nUnfortunately, this kernel won't be using a bucket and TPUs don't support reading a TFrecord on a local drive. \n\nOnce you have your own bucket, try training with files from there. You will get an error that identifies the TPU service account that is trying to access your bucket. Go to IAM on and give that TPU service account read access to that bucket or make that bucket public. \n\nYou are not in control of TPU service accounts. They are accounts that are not a part of your gcs project. "},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# for an unknown reason, the kernel dies at this point. \n# this could be because it is trying access a local tfrecord. \n\n# model.fit(dataset,\n#           epochs = 10, \n#           steps_per_epoch = 10, \n# #               validation_data = val_ds, \n# #               validation_steps = val_steps//strategy.num_replicas_in_sync,\n# #               callbacks = callbacks \n#          )\n","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"","_uuid":"","trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}