ocrorot-train

#!/usr/bin/python
import argparse

import torch
import scipy.ndimage as ndi
from pylab import *
from torch import nn, optim
from dlinputs import gopen, paths, utils, filters
from dltrainers import layers, helpers
from torch.autograd import Variable

rc("image", cmap="gray")
ion()

parser = argparse.ArgumentParser("train a page segmenter")
parser.add_argument("-l", "--lr", default="0,0.1:1e4,0.03:1e5,0.01",
                    help="learning rate or learning rate sequence 'n,lr:n,lr:n,:r'")
parser.add_argument("-b", "--batchsize", type=int, default=64)
parser.add_argument("-o", "--output", default="rot", help="prefix for output")
parser.add_argument("-m", "--model", default=None, help="load model")

parser.add_argument("--save_every", default=1000,
                    type=int, help="how often to save")
parser.add_argument("--loss_horizon", default=1000, type=int,
                    help="horizon over which to calculate the loss")
parser.add_argument("--ntrain", type=int, default=-
                    1, help="ntrain starting value")
parser.add_argument("--random_invert", type=float, default=0.0)
parser.add_argument("--min_range", type=float, default=0.4)

parser.add_argument("-D", "--makesource", default=None)
parser.add_argument("-P", "--makepipeline", default=None)
parser.add_argument("-M", "--makemodel", default=None)
parser.add_argument("--exec", dest="execute", nargs="*", default=[])
parser.add_argument(
    "--input", default="/home/tmb/tmb-ocr/uw3-patches/uw3-patches-@010.tgz")

args = parser.parse_args()
ARGS = {k: v for k, v in args.__dict__.items()}


def make_source():
    return gopen.open_source(args.input)


def make_pipeline():

    def transformer(sample):
        image = sample["input"]
        orientation = randint(0, 4)
        assert orientation < 4
        angle = 90.0 * orientation
        rotated = ndi.rotate(image, angle)
        rotated = np.expand_dims(rotated, 3)
        return dict(input=rotated, cls=orientation)

    return filters.compose(
        filters.shuffle(100, 10),
        filters.rename(input="patch.png"),
        filters.transform(transformer),
        filters.batched(args.batchsize))


def make_model():
    r = 3
    B, D, H, W = (1, 128), (1, 512), 256, 256
    model = nn.Sequential(
        layers.CheckSizes(B, D, H, W),
        nn.Conv2d(1, 8, r, padding=r//2),
        nn.BatchNorm2d(8),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),
        nn.Conv2d(8, 16, r, padding=r//2),
        nn.BatchNorm2d(16),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),
        nn.Conv2d(16, 32, r, padding=r//2),
        nn.BatchNorm2d(32),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),
        nn.Conv2d(32, 64, r, padding=r//2),
        nn.BatchNorm2d(64),
        nn.ReLU(),
        layers.Img2FlatSum(),
        nn.Linear(64, 64),
        nn.BatchNorm1d(64),
        nn.ReLU(),
        nn.Linear(64, 4),
        nn.Sigmoid(),
        layers.CheckSizes(B, 4)
    )
    return model


if args.makepipeline:
    execfile(args.makepipeline)
if args.makesource:
    execfile(args.makesource)
if args.makemodel:
    execfile(args.makemodel)
for e in args.execute:
    exec args.execute

if args.model:
    model = torch.load(args.model)
    ntrain, _ = paths.parse_save_path(args.model)
else:
    model = make_model()
    ntrain = 0

model.cuda()
source = make_source()
sample = source.next()
utils.print_sample(sample)
pipeline = make_pipeline()
source = pipeline(source)
sample = source.next()
utils.print_sample(sample)

if args.ntrain >= 0:
    ntrain = args.ntrain
print "ntrain", ntrain
print model

start_count = 0

criterion = nn.MSELoss()
criterion.cuda()

losses = [1.0]


def train_batch(model, image, cls, nclasses=4, lr=1e-3):
    cuinput = torch.FloatTensor(image.transpose(0, 3, 1, 2)).cuda()
    optimizer = optim.SGD(model.parameters(), lr=lr,
                          momentum=0.9, weight_decay=0.0)
    optimizer.zero_grad()
    cuoutput = model(Variable(cuinput))
    b, d = cuoutput.size()
    target = torch.zeros(len(cls), nclasses)
    for i, c in enumerate(cls):
        target[i, c] = 1
    cutarget = Variable(target.cuda())
    loss = criterion(cuoutput, cutarget)
    loss.backward()
    optimizer.step()
    return loss.data.cpu().numpy()[0], helpers.asnd(cuoutput)


losses = []
rates = helpers.LearningRateSchedule(args.lr)
nbatches = 0
for sample in source:
    image = sample["input"]
    cls = sample["cls"]
    lr = rates(ntrain)
    loss, output = train_batch(model, image, cls, nclasses=4, lr=lr)
    try:
        pass
    except Exception, e:
        utils.print_sample(sample)
        print e
        continue
    losses.append(loss)
    print nbatches, ntrain, loss, np.amin(output), np.amax(output), "lr", lr,
    print output[0], cls[0]
    if nbatches > 0 and nbatches % args.save_every == 0:
        err = float(np.mean(losses[-args.save_every:]))
        fname = paths.make_save_path(args.output, ntrain, err)
        torch.save(model, fname)
        print "saved", fname
    if nbatches % 10 == 0:
        clf()
        title(str(output[0]))
        imshow(image[0, :, :, 0], vmin=0, vmax=1)
        draw()
        ginput(1, 1e-3)
    waitforbuttonpress(0.0001)
    nbatches += 1
    ntrain += len(image)