code: import sys
  import os
  from os.path import expanduser
  import pickle
  import torch
  import torch.nn as nn
  import torch.optim as optim
  import torch.utils.data
  import torch.onnx
  import re
  import json
  from PIL import Image, ImageDraw
  import torch
  import numpy as np
  
  # Training script- trains a Pytorch model against the Google Quickdraw dataset:
  # https://github.com/googlecreativelab/quickdraw-dataset
  #
  # Specifically, it uses the "simplified Drawing files":
  #
  # https://console.cloud.google.com/storage/browser/quickdraw_dataset/full/simplified
  #
  # Also see https://www.kaggle.com/google/tinyquickdraw for a single downloadable tar file
  # with about 50 million samples separated into 343 classes, which is where I got mine.
  #
  # It expects those files to be in ~/data/quickdraw. Specify any alternate path on the command line.
  # 
  # As output it generates two files: doodles.pth (internal format) and doodles.onnx (ONNX export format).
  #
  # The model used here is a convolutional neural network accepting 1x64x64 inputs
  # (i.e. black-and-white 64x64 images). Output is 344 neurons (i.e. one per label) with an extra neuron
  # corresponding to label "nothing".
  # 
  # NOTES:
  # 
  # If doodles.pth is found (typically saved from a previous run), it will be loaded into the
  # current model; otherwise it will start with a set of random weights. File size is approx. 300 MB.
  # 
  # If it finds at any point during training that the output files doodles.pth or doodles.onnx
  # are not on the drive, it will write new copies immediately with its current state (even though
  # this means the first versions will only contain random weights). Deleting the files
  # generates fresh copies, and so does finishing a training epoch (overwriting the prior versions).
  # Because the data set is so immense, each epoch takes several hours to complete.
  # In practice, with this model, performance levels off after about 3-4 epochs, with the network
  # agreeing with Google's classification about 73% of the time.
  # 
  # This way, if you need to edit a hyperparameter or go to work, you can pause execution by
  # deleting the current doodles.pth and doodles.onnx files, letting it write new ones,
  # and then hitting Ctrl-C. Typically you will want to adjust the learning rate downward
  # or experiment with a different optimizer after the script has run for a few hours and
  # its performance has reached a plateau. After you make your edits the script will pick up
  # where it left off.
  #
  # If SAVE_BACKUP_FILES is set to True, the script will save backups as training progresses.
  # Each time performance reaches a new record, a file will be saved with a filename indicating the
  # new record number of correct responses. This is to avoid losing progress if the script crashes.
  # (Raising the batch size too high can cause spurious out-of-memory errors at random times.)
  
  
  # Specify data folder as command line argument; default is ~/data/quickdraw
  DATA_DIRECTORY = '~/data/quickdraw'
  if len(sys.argv) > 1:
      DATA_DIRECTORY = sys.argv[1]
  if DATA_DIRECTORY[0] == '~':
      DATA_DIRECTORY = expanduser(DATA_DIRECTORY)
  
  # Standard industry practice: Jack this number up as high as you can, then carefully lower it
  # until the script stops crashing. Final value is dependent on GPU memory.
  # This is a safe batch size to use on an RTX 2060 with 6 GB.
  BATCH_SIZE = 1000
  
  # Hyperparameters; both SGD and Adam work well, at least in the beginning; use SGD by default
  OPTIMIZER_NAME = 'SGD'
  
  SGD_LEARNING_RATE = 0.01
  SGD_MOMENTUM = 0
  
  ADAM_LEARNING_RATE = 0.001
  ADAM_BETAS = (0.9, 0.99)
  ADAM_EPSILON = 0.0001
  
  INDEX_CACHE_FILE = './index_cache.pkl'
  LABELS_FILE = './labels.txt'
  
  STATE_DICT_FILE = './doodles.pth'
  ONNX_FILE = './doodles.onnx'
  
  SAVE_BACKUP_FILES = True
  NUMBERED_STATE_DICT_FILE_TEMPLATE = './doodles_{}_of_{}.pth'
  NUMBERED_ONNX_FILE_TEMPLATE = './doodles_{}_of_{}.onnx'
  
  DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  # If it's installed, turn this on to enable NVidia's Apex AMP Pytorch extension.
  # This will let us do calculations in FP16 on the GPU which will save memory on the card
  # and let us raise the batch size. It will also leverage RTX tensor cores on RTX cards.
  # Default is set to False, because compiling and installing AMP is an involved process-
  # NVidia's CUDA Toolkit to be installed on your system before you can compile it using pip.
  
  MIXED_PRECISION = False
  
  if MIXED_PRECISION and torch.cuda.is_available():
      # See if the AMP Pytorch extension has been installed; otherwise stick to standard FP32.
      # If we are using mixed precision we can raise the batch size but keep it a multiple of 8.
      # All tensor dimensions must be multiples of 8 to trigger NVidia's tensor core optimizations.
      try:
          from apex import amp, optimizers
          MIXED_PRECISION = True
          BATCH_SIZE = int(BATCH_SIZE * 1.6) # Raising it by 60%
          print('Using mixed precision.')
      except ImportError:
          MIXED_PRECISION = False
  
  # This is a torch DataSet implementation that makes the following assumptions:
  #
  # 1. Data consists of a set of text files with ".ndjson" extensions in the specified directory.
  # 2. Each line in the .ndjson file is a JSON string with all data for a single sample.
  # 3. Each line of JSON has the following format (omitting extraneous fields):
  #    {"word":"elephant","drawing":[[[0, 1, 10],[25, 103, 163]],[[4,15,134,234,250],[27,22,6,4,0]]]}
  #    Array "drawing" has the brush strokes, each stroke a pair of arrays with x and y coordinates on a 256x256 grid.
  # 4. We can build our label list by only looking at the first line of each file. (All lines have same value for "word".)
  class QuickDrawDataset(torch.utils.data.Dataset):
  
      # Take the batch size, so we know how much to pad with all-zero samples mapping to the "blank" channel.
      # This way we ensure we deliver full-sized batches interspersed with a few blank samples mapping to label "nothing".
      def __init__(self, dataDir, batch_size):
          super(QuickDrawDataset, self).__init__()
          print('Data folder: ' + dataDir)
          self.dataDir = dataDir
          self.filenames = list(filter(lambda x: x.endswith(".ndjson"), sorted(os.listdir(dataDir)))) #[1:20]
          self.filenameByIndex = []
          self.fileByteOffsetByIndex = []
          self.labelListIndices = {}
          self.labelList = []
  
          for filename in self.filenames:
              print('Indexing ' + filename)
              file = open(dataDir + "/" + filename, "r")
              byte_offset = 0
              word = None
              for line in file:
                  if (word == None):
                      words = re.findall('\"word\":\"([\w\s-]+)\"', line)
                      word = words[0]
                      self.labelListIndices[word] = len(self.labelList)
                      self.labelList.append(word)
                  # Only use the ones Google recognizes
                  if (len(re.findall('\"recognized\":true', line)) > 0):
                      self.filenameByIndex.append(filename)
                      self.fileByteOffsetByIndex.append(byte_offset)
                  byte_offset += len(line)
              file.close()
  
          self.labelListIndices['nothing'] = len(self.labelList)
          self.labelList.append('nothing')
          if MIXED_PRECISION:
            # NVidia really wants tensor dimensions to be multiples of 8, make sure here
            extra_nothings = 0
            while len(self.labelList) % 8 > 0:
              extra_nothings += 1
              self.labelListIndices['nothing_{}'.format(extra_nothings)] = len(self.labelList)
              self.labelList.append('nothing_{}'.format(extra_nothings))
  
          self.paddingLength = batch_size - (len(self.filenameByIndex) % batch_size)
          print('padding length {}'.format(self.paddingLength))
  
      def __len__(self):
          return len(self.filenameByIndex) + self.paddingLength
  
      def __getitem__(self, idx):
          if idx >= len(self.filenameByIndex):
              # NULL sample
              return torch.zeros(1, 64, 64, dtype=torch.float), self.labelListIndices['nothing']
          filename = self.filenameByIndex[idx]
          byte_offset = self.fileByteOffsetByIndex[idx]
          file = open(self.dataDir + '/' + filename, 'r')
          file.seek(byte_offset)
          line = file.readline()
          file.close()
          # Convert line containing brush stroke coordinate list to a 256x256 image tensor using PIL
          entry = json.loads(line)
          drawing = entry.get('drawing')
          im = Image.new("L", (256, 256))
          draw = ImageDraw.Draw(im)
          for stroke in drawing:
              x_coords = stroke[0]
              y_coords = stroke[1]
              for i in range(len(x_coords) - 1):
                  draw.line((x_coords[i], y_coords[i], x_coords[i + 1], y_coords[i + 1]), fill=255, width=5)
          im = im.resize((64, 64), Image.ANTIALIAS)
          word = entry.get('word')
          imageTensor = torch.tensor(np.array(im) / 256, dtype=torch.float)
  
          # Alter image slightly to look like the inputs we're eventually going to get from the client.
          # This is a limitation imposed by JavaScript which implements "antialiasing" on downsized canvases by
          # nearest-neighbor downsampling, smoothed onscreen by a WebGL filter that looks nice but doesn't alter the image data,
          # so we only get two-color jagged images.
          #
          # Tedious workarounds are possible: https://stackoverflow.com/questions/2303690/resizing-an-image-in-an-html5-canvas
          THRESHOLD = 0.1
          imageTensor[imageTensor >= THRESHOLD] = 1.0
          imageTensor[imageTensor < THRESHOLD] = 0.0
  
          imageTensor = imageTensor.unsqueeze(0)
  
          return imageTensor, self.labelListIndices.get(word)
  
  # Takes input of size Nx1x64x64, a batch of N black and white 64x64 images.
  # Applies two convolutional layers and three fully connected layers.
  
  class CNNModel(nn.Module):
  
      # input_size is 64 (input samples are 64x64 images); num_classes is 344
      def __init__(self, input_size, num_classes):
          super(CNNModel, self).__init__()
          self.layer1 = nn.Sequential(
              nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2, bias=False),
              nn.BatchNorm2d(32),
              nn.ReLU(inplace=True),
              nn.MaxPool2d(kernel_size=2, stride=2))
          self.layer2 = nn.Sequential(
              nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2, bias=False),
              nn.BatchNorm2d(64),
              nn.ReLU(inplace=True),
              nn.MaxPool2d(kernel_size=2, stride=2))
          dimension = int(64 * pow(input_size / 4, 2))
          self.fc1 = nn.Sequential(nn.Linear(dimension, int(dimension / 4)), nn.Dropout(0.25))
          self.fc2 = nn.Sequential(nn.Linear(int(dimension / 4), int(dimension / 8)), nn.Dropout(0.25))
          self.fc3 = nn.Sequential(nn.Linear(int(dimension / 8), num_classes))
  
      def forward(self, x):
          out = self.layer1(x)
          out = self.layer2(out)
          out = out.view(out.size(0), -1)
          out = self.fc1(out)
          out = self.fc2(out)
          out = self.fc3(out)
          return out
  
  # Main part
  if __name__ == '__main__':
  
      if os.path.isfile(INDEX_CACHE_FILE):
          print("Loading {}".format(INDEX_CACHE_FILE))
          infile = open(INDEX_CACHE_FILE, 'rb')
          dataSet = pickle.load(infile)
          infile.close()
      else:
          dataSet = QuickDrawDataset(DATA_DIRECTORY, BATCH_SIZE)
          outfile = open(INDEX_CACHE_FILE, 'wb')
          pickle.dump(dataSet, outfile)
          outfile.close()
          print("Saved {}".format(INDEX_CACHE_FILE))
  
      if (os.path.isfile(LABELS_FILE) == False):
          with open(LABELS_FILE, 'w') as f:
              for label in dataSet.labelList:
                  f.write("%s\n" % label)
              f.close()
          print("Saved {}".format(LABELS_FILE))
  
      print('Total number of labels: {}'.format(len(dataSet.labelList)))
      print('Total number of samples: {}'.format(len(dataSet)))
  
      randomSampler = torch.utils.data.RandomSampler(dataSet)
      dataLoader = torch.utils.data.DataLoader(dataSet, batch_size = BATCH_SIZE, sampler = randomSampler, num_workers=4, pin_memory=True)
  
      model = CNNModel(input_size=64, num_classes=len(dataSet.labelList)).to(DEVICE)
  
      if (os.path.isfile(STATE_DICT_FILE)):
          # We found an existing doodles.pth file! Instead of starting from scratch we'll load this one.
          # and continue training it.
          print("Loading {}".format(STATE_DICT_FILE))
          state_dict = torch.load(STATE_DICT_FILE)
          model.load_state_dict(state_dict)
  
      optimizer = None
      if (OPTIMIZER_NAME == 'SGD'):
          optimizer = optim.SGD(model.parameters(), lr = SGD_LEARNING_RATE, momentum=SGD_MOMENTUM)
          print('Using SGD with learning rate {} and momentum {}'.format(SGD_LEARNING_RATE, SGD_MOMENTUM))
      elif (OPTIMIZER_NAME == 'Adam'):
          if MIXED_PRECISION:
              optimizer = optim.Adam(model.parameters(), lr = ADAM_LEARNING_RATE, betas = ADAM_BETAS, eps = ADAM_EPSILON)
          else:
              optimizer = optim.Adam(model.parameters(), lr = ADAM_LEARNING_RATE)
          print('Using Adam with learning rate {}'.format(ADAM_LEARNING_RATE))
      else:
          print('No optimizer specified!')
  
      if MIXED_PRECISION:
          # Using NVidia's AMP Pytorch extension
          model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
  
      criterion = nn.CrossEntropyLoss()
  
      ROLLING_AVERAGE_RUN_LENGTH = 100
      rolling = np.zeros(0)
      record_rolling_average = 0
      count = 0
  
      # On my computer each epoch takes about 4 hours; the script consumes ~250 watts or about 1 kWh per epoch.
      # Performance reaches a plateau after 3-4 epochs.
      for epoch in range(4):
          print('Epoch: {}'.format(epoch))
          batch_number = 0
          for i, (images, labels) in enumerate(dataLoader):
              count = count + 1
              images = images.to(DEVICE)
              labels = labels.to(DEVICE)
              optimizer.zero_grad()
              outputs = model(images)
              _, predicted = torch.max(outputs.data, 1)
              correct = (predicted == labels).sum().item()
              if (count < ROLLING_AVERAGE_RUN_LENGTH):
                  rolling = np.insert(rolling, 0, correct)
              else:
                  rolling = np.roll(rolling, 1)
                  rolling[0] = correct
              rolling_average = int(np.mean(rolling))
              loss = criterion(outputs, labels)
              if MIXED_PRECISION:
                  # Use of FP16 requires loss scaling, due to underflow error.
                  # See https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/
                  with amp.scale_loss(loss, optimizer) as scaled_loss:
                      scaled_loss.backward()
              else:
                  loss.backward()
              optimizer.step()
              print('EPOCH: {}  BATCH: {}  SIZE: {}  CORRECT: {}  (ROLLING AVG: {})'.format(epoch, batch_number, BATCH_SIZE, correct, rolling_average))
              batch_number += 1
              # print(loss.item())
  
              # To be safe, save model whenever performance reaches a new high
              if (count < 2 * ROLLING_AVERAGE_RUN_LENGTH):  # (once rolling average has had time to stabilize)
                  record_rolling_average = max(rolling_average, record_rolling_average)
              else:
                  if (rolling_average > record_rolling_average):
                      # Save model with a munged filename; e.g. doodles_706.pth
                      if (SAVE_BACKUP_FILES):
                          backupPth = NUMBERED_STATE_DICT_FILE_TEMPLATE.format(rolling_average, BATCH_SIZE)
                          torch.save(model.state_dict(), backupPth)
                          print('Saved model file {}'.format(backupPth))
                          # Delete the last backup .pth file we wrote to avoid filling up the drive
                          if (record_rolling_average > 0):
                              old_file = NUMBERED_STATE_DICT_FILE_TEMPLATE.format(record_rolling_average, BATCH_SIZE)
                              if os.path.exists(old_file):
                                  os.remove(old_file)
                          # Same for ONNX
                          backupOnnx = NUMBERED_ONNX_FILE_TEMPLATE.format(rolling_average, BATCH_SIZE)
                          if MIXED_PRECISION:
                              with amp.disable_casts():
                                  dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
                                  torch.onnx.export(model, dummy_input, backupOnnx, verbose=False)
                          else:
                              dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
                              torch.onnx.export(model, dummy_input, backupOnnx, verbose=False)
                          print('Saved ONNX file {}'.format(backupOnnx))
                          # Delete the last backup ONNX file we wrote to avoid filling up the drive
                          if (record_rolling_average > 0):
                              old_file = NUMBERED_ONNX_FILE_TEMPLATE.format(record_rolling_average, BATCH_SIZE)
                              if os.path.exists(old_file):
                                  os.remove(old_file)
                      record_rolling_average = rolling_average
  
              # Deleting the model file during training triggers a fresh rewrite:
              if (os.path.isfile(STATE_DICT_FILE) == False):
                  torch.save(model.state_dict(), STATE_DICT_FILE)
                  print('Saved model file {}'.format(STATE_DICT_FILE))
              # ONNX: same policy
              if (os.path.isfile(ONNX_FILE) == False):
                  if MIXED_PRECISION:
                      with amp.disable_casts():
                          dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
                          torch.onnx.export(model, dummy_input, ONNX_FILE, verbose=False)
                  else:
                      dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
                      torch.onnx.export(model, dummy_input, ONNX_FILE, verbose=False)
                  print('Exported ONNX file {}'.format(ONNX_FILE))
          # Epoch finished
          # Save the current model at the end of an epoch
          torch.save(model.state_dict(), STATE_DICT_FILE)
          # Export ONNX with loudmouth flag set
          if (MIXED_PRECISION):
              with amp.disable_casts():
                  dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
                  torch.onnx.export(model, dummy_input, ONNX_FILE, verbose=True)
          else:
              dummy_input = torch.randn(1, 1, 64, 64).to(DEVICE)
              torch.onnx.export(model, dummy_input, ONNX_FILE, verbose=True)
          print('EPOCH {} FINISHED, SAVED {} AND {}'.format(epoch, STATE_DICT_FILE, ONNX_FILE))
  
apis: torch.utils.data.DataLoader numpy.insert torch.cuda.is_available torch.nn.Conv2d torch.max torch.nn.Dropout torch.nn.BatchNorm2d torch.onnx.export torch.randn torch.utils.data.RandomSampler numpy.mean torch.nn.MaxPool2d torch.load numpy.zeros torch.zeros numpy.roll torch.nn.CrossEntropyLoss numpy.array torch.nn.ReLU