{"cells":[{"metadata":{},"cell_type":"markdown","source":"# How To Prepare A Stratified Split\n\nI've seen many public kernels doing a very simple split of the training data.\n\nBut because we are dealing with data from different sources, we might want to make sure that each source is represented evenly in our training and validation data.\n\nWhat follows is a quick snippet on how to do that. It's really just a few lines (without all the comments I've added for you).\n\n<p style=\"color:red\">If you find this notebook useful, please... do whatever you feel like. It's your life.</p>"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nfrom sklearn.model_selection import StratifiedKFold\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom itertools import islice\n\n# get train csv as pandas dataframe\ndf = pd.read_csv(os.path.join('../input/global-wheat-detection', 'train.csv'))\n\n# get a df with just image and source columns\n# such that dropping duplicates will only keep unique image_ids\nimage_source = df[['image_id', 'source']].drop_duplicates()\n\n# get lists for image_ids and sources\nimage_ids = image_source['image_id'].to_numpy()\nsources = image_source['source'].to_numpy()\n\n# do the split\n# in other words:\n# split up our data into 10 buckets making sure that each bucket\n#  has a more or less even distribution of sources\n# Note the use of random_state=1 to ensure the split is the same each time we run this code\nskf = StratifiedKFold(n_splits=10, shuffle=True, random_state=1)\nsplit = skf.split(image_ids, sources) # second arguement is what we are stratifying by\n\n# so now `split` is an iterator\n# each iteration gives us a set of indices pointing to the train rows of the df \n#  (in this case 90% of all the data)\n#  and a set of indices pointing to the val rows of the df\n#  (in this case 10% of the data)\n# we can use islice to control which split we select\nselect = 0\ntrain_ix, val_ix = next(islice(split, select, select+1))\n\n# translate indices to ids\ntrain_ids = image_ids[train_ix]\nval_ids = image_ids[val_ix]\n\n# create corresponding dfs\ntrain_df = df[df['image_id'].isin(train_ids)]\nval_df = df[df['image_id'].isin(val_ids)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Now we can plot the distributions to check that they are even"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"print(f'# train images: {train_ids.shape[0]}')\nprint(f'# val images: {val_ids.shape[0]}')\n\nfig = plt.figure(figsize=(20, 5))\ncounts = train_df['source'].value_counts()\nax1 = fig.add_subplot(1,2,1)\na = ax1.bar(counts.index, counts)\ncounts = val_df['source'].value_counts()\nax2 = fig.add_subplot(1,2,2)\na = ax2.bar(counts.index, counts)","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":4}