text: import os
  import argparse
  import torch
  import numpy as np
  from torch import optim
  from torch import multiprocessing
  from torch.optim import lr_scheduler
  from torch.autograd import Variable
  from torch.utils.data import DataLoader, ConcatDataset
  from util import (SingleViewTripletBuilder, distance, Logger, ensure_folder)
  from tcn import define_model
  
  IMAGE_SIZE = (299, 299)
  
  def get_args():
      parser = argparse.ArgumentParser()
      parser.add_argument('--start-epoch', type=int, default=1)
      parser.add_argument('--epochs', type=int, default=10000)
      parser.add_argument('--save-every', type=int, default=25)
      parser.add_argument('--model-folder', type=str, default='./trained_models/tcn/')
      parser.add_argument('--load-model', type=str, required=False)
      parser.add_argument('--train-directory', type=str, default='./data/train/')
      parser.add_argument('--validation-directory', type=str, default='./data/validation/')
      parser.add_argument('--minibatch-size', type=int, default=256)
      parser.add_argument('--margin', type=float, default=2.0)
      parser.add_argument('--model-name', type=str, default='tcn')
      parser.add_argument('--log-file', type=str, default='./out.log')
      parser.add_argument('--lr-start', type=float, default=0.01)
      parser.add_argument('--triplets-from-videos', type=int, default=5)
      return parser.parse_args()
  
  arguments = get_args()
  
  logger = Logger(arguments.log_file)
  def batch_size(epoch, max_size):
      exponent = epoch // 100
      return min(max(2 ** (exponent), 2), max_size)
  
  validation_builder = SingleViewTripletBuilder(arguments.validation_directory, IMAGE_SIZE, arguments, sample_size=100)
  validation_set = [validation_builder.build_set() for i in range(10)]
  validation_set = ConcatDataset(validation_set)
  del validation_builder
  
  def validate(tcn, use_cuda, arguments):
      # Run model on validation data and log results
      data_loader = DataLoader(validation_set, batch_size=256, shuffle=False, pin_memory=use_cuda)
      correct_with_margin = 0
      correct_without_margin = 0
      for minibatch, _ in data_loader:
          frames = Variable(minibatch, volatile=True)
  
          if use_cuda:
              frames = frames.cuda()
  
          anchor_frames = frames[:, 0, :, :, :]
          positive_frames = frames[:, 1, :, :, :]
          negative_frames = frames[:, 2, :, :, :]
  
          anchor_output = tcn(anchor_frames)
          positive_output = tcn(positive_frames)
          negative_output = tcn(negative_frames)
  
          d_positive = distance(anchor_output, positive_output)
          d_negative = distance(anchor_output, negative_output)
  
          assert(d_positive.size()[0] == minibatch.size()[0])
  
          correct_with_margin += ((d_positive + arguments.margin) < d_negative).data.cpu().numpy().sum()
          correct_without_margin += (d_positive < d_negative).data.cpu().numpy().sum()
  
      message = "Validation score correct with margin {with_margin}/{total} and without margin {without_margin}/{total}".format(
          with_margin=correct_with_margin,
          without_margin=correct_without_margin,
          total=len(validation_set)
      )
      logger.info(message)
  
  def model_filename(model_name, epoch):
      return "{model_name}-epoch-{epoch}.pk".format(model_name=model_name, epoch=epoch)
  
  def save_model(model, filename, model_folder):
      ensure_folder(model_folder)
      model_path = os.path.join(model_folder, filename)
      torch.save(model.state_dict(), model_path)
  
  
  def build_set(queue, triplet_builder, log):
      while 1:
          datasets = []
          for i in range(5):
              dataset = triplet_builder.build_set()
              datasets.append(dataset)
          dataset = ConcatDataset(datasets)
          log.info('Created {0} triplets'.format(len(dataset)))
          queue.put(dataset)
  
  def create_model(use_cuda):
      tcn = define_model(use_cuda)
      # tcn = PosNet()
      if arguments.load_model:
          model_path = os.path.join(
              arguments.model_folder,
              arguments.load_model
          )
          # map_location allows us to load models trained on cuda to cpu.
          tcn.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
  
      if use_cuda:
          tcn = tcn.cuda()
      return tcn
  
  
  def main():
      use_cuda = torch.cuda.is_available()
  
      tcn = create_model(use_cuda)
  
      triplet_builder = SingleViewTripletBuilder(arguments.train_directory, IMAGE_SIZE, arguments, sample_size=200)
  
      queue = multiprocessing.Queue(1)
      dataset_builder_process = multiprocessing.Process(target=build_set, args=(queue, triplet_builder, logger), daemon=True)
      dataset_builder_process.start()
  
      optimizer = optim.SGD(tcn.parameters(), lr=arguments.lr_start, momentum=0.9)
      # This will diminish the learning rate at the milestones.
      # 0.1, 0.01, 0.001
      learning_rate_scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[100, 500, 1000], gamma=0.1)
  
      ITERATE_OVER_TRIPLETS = 5
  
      for epoch in range(arguments.start_epoch, arguments.start_epoch + arguments.epochs):
          logger.info("Starting epoch: {0} learning rate: {1}".format(epoch,
              learning_rate_scheduler.get_lr()))
          learning_rate_scheduler.step()
  
          dataset = queue.get()
          logger.info("Got {0} triplets".format(len(dataset)))
  
          data_loader = DataLoader(
              dataset=dataset,
              batch_size=arguments.minibatch_size, # batch_size(epoch, arguments.max_minibatch_size),
              shuffle=True,
              pin_memory=use_cuda
          )
  
          if epoch % 10 == 0:
              validate(tcn, use_cuda, arguments)
  
          for _ in range(0, ITERATE_OVER_TRIPLETS):
              losses = []
              for minibatch, _ in data_loader:
                  frames = Variable(minibatch)
  
                  if use_cuda:
                      frames = frames.cuda()
  
                  anchor_frames = frames[:, 0, :, :, :]
                  positive_frames = frames[:, 1, :, :, :]
                  negative_frames = frames[:, 2, :, :, :]
  
                  anchor_output = tcn(anchor_frames)
                  positive_output = tcn(positive_frames)
                  negative_output = tcn(negative_frames)
  
                  d_positive = distance(anchor_output, positive_output)
                  d_negative = distance(anchor_output, negative_output)
                  loss = torch.clamp(arguments.margin + d_positive - d_negative, min=0.0).mean()
  
                  losses.append(loss.data.cpu().numpy()[0])
  
                  optimizer.zero_grad()
                  loss.backward()
                  optimizer.step()
  
              logger.info('loss: ', np.mean(losses))
  
          if epoch % arguments.save_every == 0 and epoch != 0:
              logger.info('Saving model.')
              save_model(tcn, model_filename(arguments.model_name, epoch), arguments.model_folder)
  
  
  
  
  
  if __name__ == '__main__':
      main()
  