colabs/image_classification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile config.yml\n",
    "task: image_classification # do not change\n",
    "base_model: google/vit-base-patch16-224 # the model to be used from hugging face hub\n",
    "project_name: autotrain-image-classification-model # the name of the project, must be unique\n",
    "log: tensorboard # do not change\n",
    "backend: local # do not change\n",
    "\n",
    "data:\n",
    "  path: data/ # the path to the data folder, this folder consists of `train` and `valid` (if any) folders\n",
    "  train_split: train # this folder inside data/ will be used for training, it contains the images in subfolders.\n",
    "  valid_split: null # this folder inside data/ will be used for validation, it contains the images in subfolders. If not available, set it to null\n",
    "  column_mapping: # do not change\n",
    "    image_column: image\n",
    "    target_column: labels\n",
    "\n",
    "params:\n",
    "  epochs: 2\n",
    "  batch_size: 4\n",
    "  lr: 2e-5\n",
    "  optimizer: adamw_torch\n",
    "  scheduler: linear\n",
    "  gradient_accumulation: 1\n",
    "  mixed_precision: fp16\n",
    "\n",
    "hub:\n",
    "  username: ${HF_USERNAME} # please set HF_USERNAME in colab secrets\n",
    "  token: ${HF_TOKEN} # please set HF_TOKEN in colab secrets, must be valid hugging face write token\n",
    "  push_to_hub: true # set to true if you want to push the model to the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from google.colab import userdata\n",
    "HF_USERNAME = userdata.get('HF_USERNAME')\n",
    "HF_TOKEN = userdata.get('HF_TOKEN')\n",
    "os.environ['HF_USERNAME'] = HF_USERNAME\n",
    "\n",
    "os.environ['HF_TOKEN'] = HF_TOKEN\n",
    "!autotrain --config config.yml"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}