train_latest_snapshot.py

Go to the documentation of this file.

#! /usr/bin/env python

import os
import glob
import optparse
import md5
import time

parser = optparse.OptionParser(description="""Resume caffe training given a
                                            and a directory to find the latest
                                            snapshot  solver""",
                               add_help_option=True)
parser.add_option('--solver', type=str,
                    help="Path to solver prototxt.")
parser.add_option('--logDir', type=str, default="~/",
                    help="Prefix for training log.")


required = ["solver"]

(args, options) = parser.parse_args(args=None, values=None)

def findLatest(paths):
  if len(paths) is 0:
    return None
  mtime = lambda path: os.stat(path).st_mtime
  return next(reversed(sorted(paths, key=mtime)))


argDict = vars(args)
for key in required:
  if argDict[key] is None:
    raise Exception("Argument '%s' is required." % (key))

snapshotDirLine = [line for line in open(args.solver).readlines()
                   if "snapshot_prefix" in line][0]

prefix = snapshotDirLine.split(":")[1].strip().strip('"').strip("'")

pattern = prefix + "*_iter_*.solverstate"

snapshot = findLatest(glob.glob(pattern))

print "Found recent snapshot:%s" % snapshot

jobid = os.environ["PBS_JOBID"] if "PBS_JOBID" in os.environ \
        else "interactive_%s" % md5.new(str(time.time())).hexdigest()[:5]
logName = "train_latest_snapshot_%s.out" % jobid
logPath = os.path.join(args.logDir, logName)


command = ["caffe.bin", "--gpu", "all", "train",
           "--solver", args.solver, "--snapshot", snapshot  , ">>", logPath,
           "2>&1"]

print "Starting training... \n%s\n" % " ".join(command)

os.system(" ".join(command))