train_onet.py

# coding: utf-8
import sys
from data.DataSouce import DataSource
from data.augmentation import *
from network import *
from util.Logger import Logger
import os
import random
import time
import mxnet as mx
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import loss as gloss, nn
import time

if not os.path.exists("./log/"):
    os.mkdir("./log/")
log = Logger("./log/{}_{}.log".format(__file__.split('/')[-1],
                                      time.strftime("%Y%m%d-%H%M%S"), time.localtime), level='debug').logger

train_batch = 400
display = 100

base_lr = 0.001
momentum = 0.9
weight_decay = 0.0005

lr_steps = [30000, 50000, 60000, 70000]
lr_decay = 0.1
save_interval = 5000
max_iter = 80000
ctx = mx.gpu()

save_dir = "./models"
if not os.path.exists(save_dir):
    os.mkdir(save_dir)

root_dir = r"dataset/"
INPUT_IMAGE_SIZE = 48

MEANS = [127.5, 127.5, 127.5]
train_anno_path = []

train_anno_path += [os.path.join(root_dir, "train_faces_o/pos/image_pos")]
train_anno_path += [os.path.join(root_dir, "train_faces_o/pos/label_pos")]

train_anno_path += [os.path.join(root_dir, "train_faces_o/part/image_part")]
train_anno_path += [os.path.join(root_dir, "train_faces_o/part/label_part")]

train_anno_path += [os.path.join(root_dir, "train_faces_o/neg/image_neg")]
train_anno_path += [os.path.join(root_dir, "train_faces_o/neg/label_neg")]
start_epoch = 0


net = ONet1()
net.initialize( ctx=ctx)
net.hybridize()
trainer = gluon.Trainer(
        net.collect_params(), 'sgd',
        {'learning_rate': base_lr, 'wd': weight_decay, 'momentum': momentum})

lossFun = LossFun()
eval = Evaluate()
train_dataset = DataSource(train_anno_path, transform=Compose([
        RandomMirror(0.5), SubtractFloatMeans(
            MEANS), ToPercentCoords(), PermuteCHW()
    ]), ratio=2, image_shape=(48,48,3))

for k in range(start_epoch, max_iter + 1):

    while lr_steps and k >= lr_steps[0]:
        new_lr = trainer.learning_rate * lr_decay
        lr_steps.pop(0)
        trainer.set_learning_rate(new_lr)

    images, targets = train_dataset.getbatch(train_batch)
    images = images.as_in_context(ctx)
    targets = targets.as_in_context(ctx)
    with autograd.record():

        cls, box = net(images)

        cls = cls.reshape(cls.shape[0], cls.shape[1])
        box = box.reshape(box.shape[0], box.shape[1])
        cls_loss = lossFun.AddClsLoss(cls, targets)
        box_loss = lossFun.AddRegLoss(box, targets)
        loss = cls_loss + box_loss*3
        
    loss.backward()
    trainer.step(1)
    
    if k % 100 == 0:
        cls_auc = eval.AddClsAccuracy(cls, targets)
        reg_auc = eval.AddBoxMap(box, targets, 48, 48)
        log.info("iter: {}, cls_loss: {:.4f}, box_loss: {:.4f}, toal loss: {:.4f}, lr: {}, cls_auc: {:.4f}, reg_auc: {:.4f} ".format(
            k, cls_loss.asscalar(), box_loss.asscalar(), loss.asscalar(), trainer.learning_rate, 
            cls_auc.asscalar(), reg_auc ) )

    if k % 5000 == 0:
        net.save_parameters('./models/onet_' + str(k))