train.py

from __future__ import print_function
from __future__ import division

import argparse
import random
import torch
from PIL import Image
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable
import numpy as np
# from warpctc_pytorch import CTCLoss
from torch.nn import CTCLoss
from torch.utils.data.dataset import Dataset
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# from sklearn.model_selection import train_test_split
import os
import utils
import dataset

import models.crnn as net
import params

# parser = argparse.ArgumentParser()
# parser.add_argument('-train', '--trainroot', required=True, help='path to train dataset')
# parser.add_argument('-val', '--valroot', required=True, help='path to val dataset')
# args = parser.parse_args()

# if not os.path.exists(params.expr_dir):
#     os.makedirs(params.expr_dir)

# ensure everytime the random is the same
random.seed(params.manualSeed)
np.random.seed(params.manualSeed)
torch.manual_seed(params.manualSeed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
cudnn.benchmark = True

if torch.cuda.is_available() and not params.cuda:
    print("WARNING: You have a CUDA device, so you should probably set cuda in params.py to True")

# -----------------------------------------------
"""
In this block
    Get train and val data_loader
"""    
def loader_param():
    img_x = 32
    img_y = 128
    batch_size = 4
    return(img_x, img_y, batch_size)

img_x, img_y, batch_size = loader_param()
dataset = dataset.dataset(image_root="../IAM Dataset/words/", label_root = "../IAM Dataset/ascii/labels.txt", img_x = img_x, img_y = img_y)
train_loader, val_loader = utils.dataloader(dataset = dataset, batch_size = 16, validation_split = 0.2, shuffle_dataset = True)


# def data_loader():
#     # train
#     train_dataset = dataset.lmdbDataset(root=args.trainroot)
#     assert train_dataset
#     if not params.random_sample:
#         sampler = dataset.randomSequentialSampler(train_dataset, params.batchSize)
#     else:
#         sampler = None
#     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=params.batchSize, \
#             shuffle=True, sampler=sampler, num_workers=int(params.workers), \
#             collate_fn=dataset.alignCollate(imgH=params.imgH, imgW=params.imgW, keep_ratio=params.keep_ratio))
    
#     # val
#     val_dataset = dataset.lmdbDataset(root=args.valroot, transform=dataset.resizeNormalize((params.imgW, params.imgH)))
#     assert val_dataset
#     val_loader = torch.utils.data.DataLoader(val_dataset, shuffle=True, batch_size=params.batchSize, num_workers=int(params.workers))
    
#     return train_loader, val_loader

# train_loader, val_loader = data_loader()

# -----------------------------------------------
"""
In this block
    Net init
    Weight init
    Load pretrained model
"""
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def net_init():
    nclass = len(params.alphabet) + 1
    crnn = net.CRNN(params.imgH, params.nc, nclass, params.nh)
    crnn.apply(weights_init)
    if params.pretrained != '':
        print('loading pretrained model from %s' % params.pretrained)
        if params.multi_gpu:
            crnn = torch.nn.DataParallel(crnn)
        crnn.load_state_dict(torch.load(params.pretrained))
    
    return crnn

# crnn = net_init()
crnn = net.CRNN(32, 3, 79, 1024)
print(crnn)

# -----------------------------------------------
"""
In this block
    Init some utils defined in utils.py
"""
# Compute average for `torch.Variable` and `torch.Tensor`.
loss_avg = utils.averager()

# Convert between str and label.
converter = utils.strLabelConverter()

# -----------------------------------------------
"""
In this block
    criterion define
"""
criterion = CTCLoss()

# -----------------------------------------------
"""
In this block
    Init some tensor
    Put tensor and net on cuda
    NOTE:
        image, text, length is used by both val and train
        becaues train and val will never use it at the same time.
"""
image = torch.FloatTensor(params.batchSize, 3, params.imgH, params.imgH)
text = torch.LongTensor(params.batchSize * 5)
length = torch.LongTensor(params.batchSize)

if params.cuda and torch.cuda.is_available():
    criterion = criterion.cuda()
    image = image.cuda()
    text = text.cuda()

    crnn = crnn.cuda()
    if params.multi_gpu:
        crnn = torch.nn.DataParallel(crnn, device_ids=range(params.ngpu))

image = Variable(image)
text = Variable(text)
length = Variable(length)

# -----------------------------------------------
"""
In this block
    Setup optimizer
"""
if params.adam:
    optimizer = optim.Adam(crnn.parameters(), lr=params.lr, betas=(params.beta1, 0.999))
elif params.adadelta:
    optimizer = optim.Adadelta(crnn.parameters())
else:
    optimizer = optim.RMSprop(crnn.parameters(), lr=params.lr)

# -----------------------------------------------
"""
In this block
    Dealwith lossnan
    NOTE:
        I use different way to dealwith loss nan according to the torch version. 
"""
if params.dealwith_lossnan:
    if torch.__version__ >= '1.1.0':
        """
        zero_infinity (bool, optional):
            Whether to zero infinite losses and the associated gradients.
            Default: ``False``
            Infinite losses mainly occur when the inputs are too short
            to be aligned to the targets.
        Pytorch add this param after v1.1.0 
        """
        criterion = CTCLoss(zero_infinity = True)
    else:
        """
        only when
            torch.__version__ < '1.1.0'
        we use this way to change the inf to zero
        """
        crnn.register_backward_hook(crnn.backward_hook)

# -----------------------------------------------

# def val(net, criterion):
#     print('Start val')

#     for p in crnn.parameters():
#         p.requires_grad = False

#     net.eval()
#     val_iter = iter(val_loader)

#     i = 0
#     n_correct = 0
#     loss_avg = utils.averager() # The blobal loss_avg is used by train

#     max_iter = len(val_loader)
#     for i in range(max_iter):
#         data = val_iter.next()
#         i += 1
#         cpu_images, cpu_texts = data
#         text, length = converter.words_rep(labels_str = cpu_texts, max_out_chars = 20, batch_size = params.batchSize, device = device)
#         batch_size = cpu_images.size(0)
#         utils.loadData(image, cpu_images)
#         # t, l = converter.encode(cpu_texts)
#         # utils.loadData(text, t)
#         # utils.loadData(length, l)

#         preds = crnn(image)
#         preds_size = Variable(torch.LongTensor([preds.size(0)] * batch_size))
#         cost = criterion(preds, text, preds_size, length) / batch_size
#         loss_avg.add(cost)

#         _, preds = preds.max(2)
#         preds = preds.transpose(1, 0).contiguous().view(-1)
#         sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
#         print(sim_preds.shape)
#         cpu_texts_decode = []
#         for i in cpu_texts:
#             cpu_texts_decode.append(i.decode('utf-8', 'strict'))
#         for pred, target in zip(sim_preds, cpu_texts_decode):
#             if pred == target:
#                 n_correct += 1

#     raw_preds = converter.decode(preds.data, preds_size.data, raw=True)[:params.n_val_disp]
#     for raw_pred, pred, gt in zip(raw_preds, sim_preds, cpu_texts_decode):
#         print('%-20s => %-20s, gt: %-20s' % (raw_pred, pred, gt))

#     accuracy = n_correct / float(max_iter * params.batchSize)
#     print('Val loss: %f, accuracy: %f' % (loss_avg.val(), accuracy))

def val(net, criterion, max_iter=100):
    print('Start val')

    for p in crnn.parameters():
        p.requires_grad = False

    net.eval()
    val_iter = iter(val_loader)

    i = 0
    n_correct = 0
    loss_avg = utils.averager()

    max_iter = len(val_loader)
    for i in range(max_iter):
        data = val_iter.next()
        i += 1
        cpu_images, cpu_texts = data
        batch_size = cpu_images.size(0)
        utils.loadData(image, cpu_images)
        # t, l = converter.encode(cpu_texts)
        # utils.loadData(text, t)
        # utils.loadData(length, l)

        preds = crnn(image)
        preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
        cost = criterion(preds, text, preds_size, length) / batch_size
        loss_avg.add(cost)

        _, preds = preds.max(2)
        preds = preds.squeeze(2)
        preds = preds.transpose(1, 0).contiguous().view(-1)
        sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
        for pred, target in zip(sim_preds, cpu_texts):
            if pred == target.lower():
                n_correct += 1

    raw_preds = converter.decode(preds.data, preds_size.data, raw=True)[:opt.n_test_disp]
    for raw_pred, pred, gt in zip(raw_preds, sim_preds, cpu_texts):
        print('%-20s => %-20s, gt: %-20s' % (raw_pred, pred, gt))

    accuracy = n_correct / float(max_iter * opt.batchSize)
    print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy))


def train(net, criterion, optimizer, train_iter):
    for p in crnn.parameters():
        p.requires_grad = True
    crnn.train()

    data = train_iter.next()
    cpu_images, cpu_texts = data
    text, length = converter.words_rep(labels_str = cpu_texts, max_out_chars = 20, batch_size = params.batchSize, device = device)
    batch_size = cpu_images.size(0)
    utils.loadData(image, cpu_images)
    # t, l = converter.encode(cpu_texts)
    # utils.loadData(text, t)
    # utils.loadData(length, l)
    
    preds = crnn(image)
    preds_size = Variable(torch.LongTensor([preds.size(0)] * batch_size))
    # print("Label: ", text[0], '\nOutput: ', torch.argmax(preds, 2)[:, 0])
    cost = criterion(preds, text, preds_size, length) / batch_size
    
    optimizer.zero_grad()
    cost.backward()
    optimizer.step()
    return cost


if __name__ == "__main__":
    for epoch in range(params.nepoch):
        train_iter = iter(train_loader)
        i = 0
        while i < len(train_loader):
            cost = train(crnn, criterion, optimizer, train_iter)
            loss_avg.add(cost)
            i += 1

            if i % params.displayInterval == 0:
                print('[%d/%d][%d/%d] Loss: %f' %
                      (epoch, params.nepoch, i, len(train_loader), loss_avg.val()))
                loss_avg.reset()

            if i % params.valInterval == 0:
                val(crnn, criterion)

            # do checkpointing
            if i % params.saveInterval == 0:
                torch.save(crnn.state_dict(), '{0}/netCRNN_{1}_{2}.pth'.format(params.expr_dir, epoch, i))