Exemple #1
0
    def __init__(self, train_from_scratch=False, eval_use_only=False):
        """
    age prediction model, provide APIs for public uses
    :param onlyTrainLastLayers: is set, freeze params of previous layers
    :param train_from_scratch: if set, do not load pretrained params
    :param eval_use_only: if set, model will not load training/testing data
    """
        self.from_scratch = train_from_scratch
        self.model = RetrainedResnetModel()
        self.use_gpu = torch.cuda.is_available()
        self.transformer = resnet_transformer()
        self.checkpoint_best = config.model + "best_torch.nn"
        self.checkpoint_last = config.model + "last_torch.nn"
        self.latest_weights = config.model + "latest_torch.nn"
        self.csv_path = config.model + "log.csv"

        self.batch_size = int(parser['TRAIN']['batch_size'])
        self.num_epochs = int(parser['TRAIN']['num_epochs'])
        self.loading_jobs = int(parser['TRAIN']['jobs_to_load_data'])
        self.max_no_reduce = int(parser['TRAIN']['max_no_reduce'])

        self.weight_decay = float(parser['TRAIN']['weight_decay'])
        self.age_divide = float(parser['DATA']['age_divide'])
        self.min_lr_rate = float(parser['TRAIN']['min_lr_rate'])
        self.lr_reduce_by = float(parser['TRAIN']['lr_reduce_by'])
        self.lr_rate = float(parser['TRAIN']['init_lr_rate'])

        self.loaded = False
        self.age_criterion = nn.L1Loss()
        self.gender_criterion = nn.NLLLoss()
        self.aligner = FaceAligner()

        if self.use_gpu:
            self.model = self.model.cuda()
            self.age_criterion = self.age_criterion.cuda()
            self.gender_criterion = self.gender_criterion.cuda()

        columns = [
            'Timstamp', 'Epoch', 'Phase', 'Gender Acc', 'Age MAE',
            'Best Gender Acc', 'Best Age MAE', 'Lr_rate'
        ]
        self.csv_checkpoint = pd.DataFrame(data=[], columns=columns)
        if not self.from_scratch and os.path.exists(self.csv_path):
            self.csv_checkpoint = pd.read_csv(self.csv_path)

        if not eval_use_only:
            self.load_data()
Exemple #2
0
import numpy as np
import time
import cv2
import utils
from align_faces import FaceAligner
from predict import FacePredict

net = cv2.dnn.readNetFromCaffe("deploy.prototxt.txt", "resnet10.caffemodel")
print("Model Loaded!")

cap = cv2.VideoCapture(0)
time.sleep(2.0)
print("Starting...")

faceAligner = FaceAligner("dlibModel.dat", 384)
userid = 1
counter = 0

# loop over the frames from the video stream
while True:
    # grab the frame from the threaded video stream and resize it
    # to have a maximum width of 400 pixels
    ret, frame = cap.read()
    frame = cv2.resize(frame, (640, 480))

    # grab the frame dimensions and convert it to a blob
    (h, w) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                 (300, 300), (104.0, 177.0, 123.0))

    # pass the blob through the network and obtain the detections and
    pwd = os.getcwd()
    if clean:
        clear_dir(config.named)
    os.chdir(config.aligned)
    for img in glob.glob("*.jpg"):
        name = re.findall(r'[^_]*_[^_]*_([\D]*)[0-9]*.jpg', img)
        if not len(name): continue
        name = name[0].lower()
        if not os.path.exists(config.named + name + '/'):
            os.mkdir(config.named + name + '/')
        copyfile(img, config.named + name + '/' + img)
    os.chdir(pwd)


# TODO: any other ways to get around this public variable?
FL = FaceAligner()


def sub_align_face(picname):
    """
  sub thread function to get and store aligned faces
  :param picname: pic names
  :return: 
  """
    aligned = FL.getAligns(picname)
    if len(aligned) == 0:
        return
        # copyfile(picname, config.aligned + picname)
    cv2.imwrite(config.aligned + picname, aligned[0])

    def __init__(self,
                 load_best=True,
                 model_name='res18_cls70',
                 eval_use_only=False,
                 new_last_layer=False,
                 new_training_process=True):
        """
    :param load_best: if set, load the best weight, else load the latest weights,
                      usually load the best when doing evaluation, and latest when
                      doing training.
    :param model_name: name used for saving model weight and training info.
    :param eval_use_only: if set, model will not load training/testing data, and 
                          change the bahavior of some layers(dropout, batch norm).
    :param new_last_layer: if the model only changed last fully connected layers,
                            if set, only train last fully connected layers at first 
                            2 epochs.
    :param new_training_process: if set, create a new model and start training.
    """
        # init params
        self.model = AgeGenPredModel()
        self.model_name = model_name
        self.use_gpu = torch.cuda.is_available()
        self.transformer = image_transformer()
        self.load_best = load_best
        self.new_train = new_training_process
        self.new_last_layer = new_last_layer
        self.checkpoint_best = config.model + "{}_best.nn".format(
            model_name.lower())
        self.checkpoint_last = config.model + "{}_last.nn".format(
            model_name.lower())
        self.csv_path = config.model + self.model_name + ".csv"

        # training details
        self.batch_size = int(parser['TRAIN']['batch_size'])
        self.num_epochs = int(parser['TRAIN']['num_epochs'])
        self.loading_jobs = int(parser['TRAIN']['jobs_to_load_data'])
        self.max_no_reduce = int(parser['TRAIN']['max_no_reduce'])
        self.age_cls_unit = int(parser['RacNet']['age_cls_unit'])
        self.weight_decay = float(parser['TRAIN']['weight_decay'])
        self.age_divide = float(parser['DATA']['age_divide'])
        self.min_lr_rate = float(parser['TRAIN']['min_lr_rate'])
        self.lr_reduce_by = float(parser['TRAIN']['lr_reduce_by'])
        self.lr_rate = float(parser['TRAIN']['init_lr_rate'])

        # reduce loss on gender so the model focus on age pred
        self.reduce_gen_loss = float(parser['TRAIN']['reduce_gen_loss'])
        self.reduce_age_mae = float(parser['TRAIN']['reduce_age_mae'])

        self.weight_loaded = False
        self.age_cls_criterion = nn.BCELoss(weight=loss_weight)
        self.age_rgs_criterion = nn.L1Loss()
        self.gender_criterion = nn.CrossEntropyLoss()
        self.aligner = FaceAligner()

        if self.use_gpu:
            self.model = self.model.cuda()
            self.age_cls_criterion = self.age_cls_criterion.cuda()
            self.age_rgs_criterion = self.age_rgs_criterion.cuda()
            self.gender_criterion = self.gender_criterion.cuda()

        # csv checkpoint details
        columns = [
            'Timstamp', 'Epoch', 'Phase', 'AGE ACC', 'AGE MAE', 'GEN ACC',
            'BEST AGE ACC', 'BEST AGE MAE', 'BEST GEN ACC', 'Lr_rate'
        ]
        self.csv_checkpoint = pd.DataFrame(data=[], columns=columns)
        if not self.new_train and os.path.exists(self.csv_path):
            self.csv_checkpoint = pd.read_csv(self.csv_path)

        # load no training data when evaluation,
        if not eval_use_only:
            self.load_data()
class AgePredModel:
    """ train/test class for age/gender prediction """
    def __init__(self,
                 load_best=True,
                 model_name='res18_cls70',
                 eval_use_only=False,
                 new_last_layer=False,
                 new_training_process=True):
        """
    :param load_best: if set, load the best weight, else load the latest weights,
                      usually load the best when doing evaluation, and latest when
                      doing training.
    :param model_name: name used for saving model weight and training info.
    :param eval_use_only: if set, model will not load training/testing data, and 
                          change the bahavior of some layers(dropout, batch norm).
    :param new_last_layer: if the model only changed last fully connected layers,
                            if set, only train last fully connected layers at first 
                            2 epochs.
    :param new_training_process: if set, create a new model and start training.
    """
        # init params
        self.model = AgeGenPredModel()
        self.model_name = model_name
        self.use_gpu = torch.cuda.is_available()
        self.transformer = image_transformer()
        self.load_best = load_best
        self.new_train = new_training_process
        self.new_last_layer = new_last_layer
        self.checkpoint_best = config.model + "{}_best.nn".format(
            model_name.lower())
        self.checkpoint_last = config.model + "{}_last.nn".format(
            model_name.lower())
        self.csv_path = config.model + self.model_name + ".csv"

        # training details
        self.batch_size = int(parser['TRAIN']['batch_size'])
        self.num_epochs = int(parser['TRAIN']['num_epochs'])
        self.loading_jobs = int(parser['TRAIN']['jobs_to_load_data'])
        self.max_no_reduce = int(parser['TRAIN']['max_no_reduce'])
        self.age_cls_unit = int(parser['RacNet']['age_cls_unit'])
        self.weight_decay = float(parser['TRAIN']['weight_decay'])
        self.age_divide = float(parser['DATA']['age_divide'])
        self.min_lr_rate = float(parser['TRAIN']['min_lr_rate'])
        self.lr_reduce_by = float(parser['TRAIN']['lr_reduce_by'])
        self.lr_rate = float(parser['TRAIN']['init_lr_rate'])

        # reduce loss on gender so the model focus on age pred
        self.reduce_gen_loss = float(parser['TRAIN']['reduce_gen_loss'])
        self.reduce_age_mae = float(parser['TRAIN']['reduce_age_mae'])

        self.weight_loaded = False
        self.age_cls_criterion = nn.BCELoss(weight=loss_weight)
        self.age_rgs_criterion = nn.L1Loss()
        self.gender_criterion = nn.CrossEntropyLoss()
        self.aligner = FaceAligner()

        if self.use_gpu:
            self.model = self.model.cuda()
            self.age_cls_criterion = self.age_cls_criterion.cuda()
            self.age_rgs_criterion = self.age_rgs_criterion.cuda()
            self.gender_criterion = self.gender_criterion.cuda()

        # csv checkpoint details
        columns = [
            'Timstamp', 'Epoch', 'Phase', 'AGE ACC', 'AGE MAE', 'GEN ACC',
            'BEST AGE ACC', 'BEST AGE MAE', 'BEST GEN ACC', 'Lr_rate'
        ]
        self.csv_checkpoint = pd.DataFrame(data=[], columns=columns)
        if not self.new_train and os.path.exists(self.csv_path):
            self.csv_checkpoint = pd.read_csv(self.csv_path)

        # load no training data when evaluation,
        if not eval_use_only:
            self.load_data()

    def load_data(self):
        """
    initiate dataloader processes
    :return: 
    """
        print("[AgePredModel] load_data: start loading...")
        image_datasets = {
            x: FaceDataset(config.pics + x + '/', self.transformer[x])
            for x in ['train', 'val']
        }
        self.dataloaders = {
            x: torch.utils.data.DataLoader(image_datasets[x],
                                           batch_size=self.batch_size,
                                           shuffle=True,
                                           num_workers=self.loading_jobs)
            for x in ['train', 'val']
        }
        self.dataset_sizes = {
            x: len(image_datasets[x])
            for x in ['train', 'val']
        }
        print(
            "[AgePredModel] load_data: Done! Get {} for train and {} for test!"
            .format(self.dataset_sizes['train'], self.dataset_sizes['val']))
        print("[AgePredModel] load_data: loading finished !")

    @staticmethod
    def rand_init_layer(m):
        """
    initialization method
    :param m: torch.module
    :return: 
    """
        if isinstance(m, nn.Conv2d):
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            size = m.weight.size()
            fan_out = size[0]  # number of rows
            fan_in = size[1]  # number of columns
            variance = np.sqrt(2.0 / (fan_in + fan_out))
            m.weight.data.normal_(0.0, variance)

    def soft_load_statedic(self, state_dict):
        """
    WARNING: Always run model = nn.DataParallel after this!
    load network parameters in a soft way, the original load_statedic
     func from torch is prone to raise exceptions when mismatch. this 
     function skip all incapatible weights and print the info intead of
     raising a exception.
    :param state_dict: saved dict 
    :return: 
    """
        # remove `module.` prefix when using nn.DataParallel
        new_state_dict = OrderedDict()
        for name, weight in state_dict.items():
            if len(name) >= 7 and name[:7].lower() == 'module.':
                name = name[7:]
            new_state_dict[name] = weight
        state_dict = new_state_dict

        # start loading
        own_state = self.model.state_dict()
        error_layers = []
        for name, param in state_dict.items():
            if name in own_state:
                if isinstance(param, nn.Parameter):
                    param = param.data
                try:
                    own_state[name].copy_(param)
                except Exception:
                    print(
                        '[soft_load_statedic] WARNING: incapatible dim found for {} = {} != {}.'
                        .format(name, own_state[name].size(), param.size()))
                    error_layers.append(name)
            else:
                print(
                    '[soft_load_statedic] Unexpected key "{}" in saved state_dict'
                    .format(name))
        missing = set.union(
            set(own_state.keys()) - set(state_dict.keys()), set(error_layers))
        if len(missing) > 0:
            print('[soft_load_statedic] keys in state_dict: "{}" not loaded!'.
                  format(missing))
        return

    def train_model(self):
        print("[AgePredModel] train_model: Start training...")

        # 1.0.0.0 define Vars
        best_gen_acc = 0.
        best_age_acc = 0.
        best_age_mae = 99.
        not_reduce_rounds = 0

        # 2.0.0.0 init optimizer
        self.optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                                           self.model.parameters()),
                                    lr=self.lr_rate,
                                    weight_decay=self.weight_decay)

        # 3.0.0.0 load weights if possible
        checkpoint_path = self.checkpoint_best if self.load_best else self.checkpoint_last
        if self.new_train:
            print("[new_training_process] NO WEIGHT LOADED!")
        elif os.path.exists(checkpoint_path):
            checkpoint = torch.load(
                checkpoint_path, map_location=None if self.use_gpu else 'cpu')
            self.soft_load_statedic(checkpoint['state_dic'])
            print("[train_model] Params Loading Finished!")
            self.weight_loaded = True
            try:
                best_gen_acc = checkpoint['best_gen_acc']
                best_age_acc = checkpoint['best_age_acc']
                best_age_mae = checkpoint['best_age_mae']
                # self.lr_rate = checkpoint['lr_rate']
                self.optimizer.load_state_dict(checkpoint['optimizer'])
                for param_group in self.optimizer.param_groups:
                    param_group['lr'] = self.lr_rate
                print("[train_model] Load Optimizer Successful!")
            except:
                print("[train_model] ERROR: Loading Params/Optimizer Error!")
        else:
            print("[train_model] Checkpoint Not Found, Train From Scratch!")

        # report model params
        all_params = sum([np.prod(p.size()) for p in self.model.parameters()])
        trainable_params = sum([
            np.prod(p.size())
            for p in filter(lambda p: p.requires_grad, self.model.parameters())
        ])
        print(
            "[AgePredModel] Model has {}k out of {}k trainable params ".format(
                trainable_params // 1000, all_params // 1000))

        # use when having multiple GPUs available
        if torch.cuda.device_count() > 1:
            self.model = nn.DataParallel(self.model)

        # 4.0.0.0 start each epoch
        layer_to_freeze = 0
        for epoch in range(self.num_epochs):
            print('\nStart Epoch {}/{} ...'.format(epoch + 1, self.num_epochs))
            print('-' * 16)

            # automatically freeze some layers on first 2 epochs
            if epoch == 0:
                new_layer_to_freeze = 8  # resnet-18 has 8 modules in pytorch
            elif epoch == 1:
                new_layer_to_freeze = 6
            else:
                new_layer_to_freeze = 0
            if (self.new_last_layer or self.new_train) \
                    and layer_to_freeze != new_layer_to_freeze:
                layer_to_freeze = new_layer_to_freeze
                # free some layers
                model = self.model
                if torch.cuda.device_count() > 1:
                    model = self.model.module
                for i, child in enumerate(model.resNet.children()):
                    requires_grad = i >= int(layer_to_freeze)
                    for param in child.parameters():
                        param.requires_grad = requires_grad
                # re-define the optimizer
                self.optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                                                   model.parameters()),
                                            lr=self.lr_rate,
                                            weight_decay=self.weight_decay)

            # 4.1.0.0 loop over training and validation phase
            for phase in ['train', 'val']:
                # 4.1.1.0 shift train/eval model
                self.model.train(phase == 'train')
                torch.cuda.empty_cache()

                epoch_age_tp = 0.
                epoch_age_mae = 0.
                epoch_gender_tp = 0.
                processed_data = 0

                # 4.1.2.0 iterate over each batch.
                epoch_start_time = time.time()
                for data in self.dataloaders[phase]:
                    # 4.1.2.1 get the inputs and labels
                    inputs, attention_area, gender_true, age_rgs_true, age_cls_true = data
                    processed_data += inputs.size(0)

                    # 4.1.2.2 wrap inputs&oputpus into Variable
                    #         NOTE: set voloatile = True when
                    #         doing evaluation helps reduce
                    #         gpu mem usage.
                    volatile = phase == 'val'
                    if self.use_gpu:
                        inputs = Variable(inputs.cuda(), volatile=volatile)
                        gender_true = Variable(gender_true.cuda(),
                                               volatile=volatile)
                        # age_rgs_true  = Variable(age_rgs_true.cuda(), volatile=volatile)
                        age_cls_true = Variable(age_cls_true.cuda(),
                                                volatile=volatile)
                    else:
                        inputs = Variable(inputs, volatile=volatile)
                        gender_true = Variable(gender_true, volatile=volatile)
                        # age_rgs_true  = Variable(age_rgs_true, volatile=volatile)
                        age_cls_true = Variable(age_cls_true,
                                                volatile=volatile)

                    # 4.1.2.3 zero gradients
                    self.optimizer.zero_grad()

                    # 4.1.2.4 forward and get outputs
                    gender_out, age_out = self.model(inputs, attention_area)
                    _, gender_pred = torch.max(gender_out, 1)
                    _, max_cls_pred_age = torch.max(age_out, 1)
                    gender_true = gender_true.view(-1)
                    age_cls_true = age_cls_true.view(-1, self.age_cls_unit)

                    # 4.1.2.5 get the loss
                    gender_loss = self.gender_criterion(
                        gender_out, gender_true)
                    age_cls_loss = self.age_cls_criterion(
                        age_out, age_cls_true)
                    # age_rgs_loss  = self.age_rgs_criterion(age_out, age_rgs_true)

                    # *Note: reduce some age loss and gender loss
                    #         enforce the model to focuse on reducing
                    #         age classification loss
                    gender_loss *= self.reduce_gen_loss
                    # age_rgs_loss *= self.reduce_age_mae

                    # loss = gender_loss + age_rgs_loss + age_cls_loss
                    # loss = age_rgs_loss
                    loss = age_cls_loss
                    loss = gender_loss + age_cls_loss

                    gender_loss_perc = 100 * (gender_loss /
                                              loss).cpu().data.numpy()
                    age_cls_loss_perc = 100 * (age_cls_loss /
                                               loss).cpu().data.numpy()
                    # age_rgs_loss_perc = 100 * (age_rgs_loss / loss).cpu().data.numpy()[0]

                    age_rgs_loss_perc = 0
                    # age_cls_loss_perc = 0
                    # gender_loss_perc = 0

                    # convert cls result to rgs result by weigted sum
                    weigh = np.linspace(1, self.age_cls_unit,
                                        self.age_cls_unit)
                    age_cls_raw = age_out.cpu().data.numpy()
                    age_cls_raw = np.sum(age_cls_raw * weigh, axis=1)
                    age_rgs_true = age_rgs_true.view(-1)
                    age_rgs_true = age_rgs_true.cpu().numpy() * self.age_divide
                    age_rgs_loss = np.mean(np.abs(age_cls_raw - age_rgs_true))

                    # 4.1.2.6 backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        self.optimizer.step()

                    # 4.1.2.7 statistics
                    gender_pred = gender_pred.cpu().data.numpy()
                    gender_true = gender_true.cpu().data.numpy()
                    batch_gender_tp = np.sum(gender_pred == gender_true)

                    max_cls_pred_age = max_cls_pred_age.cpu().data.numpy()
                    age_cls_true = age_rgs_true
                    batch_age_tp = np.sum(
                        np.abs(age_cls_true - max_cls_pred_age) <=
                        2)  # if true, MAE < 5

                    epoch_age_mae += age_rgs_loss * inputs.size(0)
                    epoch_age_tp += batch_age_tp
                    epoch_gender_tp += batch_gender_tp

                    # 4.1.2.8 print info for each bach done
                    print(
                        "|| {:.2f}% {}/{} || LOSS = {:.2f} || DISTR% {:.0f} : {:.0f} : {:.0f} "
                        "|| AMAE/AACC±2/GACC = {:.2f} / {:.2f}% / {:.2f}% "
                        "|| LR {} || ETA {:.0f}s || BEST {:.2f} / {:.2f}% / {:.2f}% ||"
                        .format(
                            100 * processed_data / self.dataset_sizes[phase],
                            processed_data, self.dataset_sizes[phase],
                            loss.cpu().data.numpy(), age_rgs_loss_perc,
                            age_cls_loss_perc, gender_loss_perc, age_rgs_loss,
                            100 * batch_age_tp / inputs.size(0),
                            100 * batch_gender_tp / inputs.size(0),
                            self.lr_rate,
                            (self.dataset_sizes[phase] - processed_data) *
                            (time.time() - epoch_start_time) / processed_data,
                            best_age_mae, 100 * best_age_acc,
                            100 * best_gen_acc),
                        end='\r')

                    # 4.1.2.9 unlink cuda variables and free up mem
                    del inputs, gender_true, age_rgs_true, age_cls_true
                    del age_rgs_loss, loss  # , gen_loss, age_cls_loss
                    del gender_loss_perc, age_cls_loss_perc, age_rgs_loss_perc

                # 4.1.3.0 epoch done
                epoch_gender_acc = epoch_gender_tp / self.dataset_sizes[phase]
                epoch_age_acc = epoch_age_tp / self.dataset_sizes[phase]
                epoch_age_mae /= self.dataset_sizes[phase]

                # 4.1.4.0 print info after each epoch done
                print('\n--{} {}/{} Done! '
                      '|| AMAE/AACC±2/GACC = {:.2f} / {:.2f}% / {:.2f}%  '
                      '|| COST {:.0f}s'.format(phase.upper(), epoch,
                                               self.num_epochs, epoch_age_mae,
                                               100 * epoch_age_acc,
                                               100 * epoch_gender_acc,
                                               time.time() - epoch_start_time))

                # 4.1.5.0, save model weights
                if phase == 'val' and epoch_age_mae < best_age_mae:
                    best_gen_acc = epoch_gender_acc
                    best_age_acc = epoch_age_acc
                    best_age_mae = epoch_age_mae
                    best_model_wts = copy.deepcopy(self.model.state_dict())
                    torch.save(
                        {
                            'epoch': epoch,
                            'state_dic': best_model_wts,
                            "best_gen_acc": best_gen_acc,
                            "best_age_acc": best_age_acc,
                            "best_age_mae": best_age_mae,
                            "lr_rate": self.lr_rate,
                            "optimizer": self.optimizer.state_dict()
                        }, self.checkpoint_best)
                    not_reduce_rounds = 0
                    print(
                        "--New BEST FOUND!! || "
                        " AMAE/AACC/AACC±2/GACC = {:.2f} / {:.2f}% / {:.2f}%".
                        format(best_age_mae, 100 * best_age_acc,
                               100 * best_gen_acc))
                elif phase == 'val':
                    not_reduce_rounds += 1
                    torch.save(
                        {
                            'epoch': epoch,
                            'state_dic': self.model.state_dict(),
                            "best_gen_acc": best_gen_acc,
                            "best_age_acc": best_age_acc,
                            "best_age_mae": best_age_mae,
                            "lr_rate": self.lr_rate,
                            "optimizer": self.optimizer.state_dict()
                        }, self.checkpoint_last)

                # 4.1.6.0 save csv logging file
                try:
                    self.csv_checkpoint.loc[len(self.csv_checkpoint)] = [
                        str(datetime.datetime.now()), epoch, phase,
                        epoch_age_acc, epoch_age_mae, epoch_gender_acc,
                        best_age_acc, best_age_mae, best_gen_acc, self.lr_rate
                    ]
                    self.csv_checkpoint.to_csv(self.csv_path, index=False)
                except:
                    print(
                        "Error when saving csv files! [tip]: Please check csv column names."
                    )
                    print(self.csv_checkpoint.columns)

                # 4.1.7.0 reduce learning rate if nessessary
                if phase == "val" \
                        and not_reduce_rounds >= self.max_no_reduce \
                        and self.lr_rate > self.min_lr_rate:
                    self.lr_rate = max(self.min_lr_rate,
                                       self.lr_rate / self.lr_reduce_by)
                    print(
                        "[reduce_lr_rate] Reduce Learning Rate From {} --> {}".
                        format(self.lr_rate * self.lr_reduce_by, self.lr_rate))
                    for param_group in self.optimizer.param_groups:
                        param_group['lr'] = self.lr_rate
                    not_reduce_rounds = 0

                    # 4.2.0.0 train/val loop ends
        # 5.0.0.0 Trainning Completes!
        return self.model


# """
#
#   # evaluate function is just a pruned version of train function
#
#   def evaluate(self):
#     checkpoint_path = self.checkpoint_best if self.load_best else self.checkpoint_last
#     checkpoint = torch.load(checkpoint_path, map_location=None if self.use_gpu else 'cpu')
#     self.soft_load_statedic(checkpoint['state_dic'])
#     self.model.train(mode=False)
#
#     epoch_age_tp = 0.
#     epoch_age_mae = 0.
#     epoch_gender_tp = 0.
#     processed_data = 0
#
#     # 4.1.2.0 Iterate over data.
#     epoch_start_time = time.time()
#     phase = 'val'
#     for data in self.dataloaders[phase]:
#       # 4.1.2.1 get the inputs and labels
#       inputs, gender_true, age_rgs_true, age_cls_true = data
#       processed_data += self.batch_size
#
#       # 4.1.2.2 wrap inputs&oputpus into Variable
#       #         NOTE: set voloatile = True when
#       #         doing evaluation helps reduce
#       #         gpu mem usage.
#       volatile = phase == 'val'
#       if self.use_gpu:
#         inputs = Variable(inputs.cuda(), volatile=volatile)
#         gender_true = Variable(gender_true.cuda(), volatile=volatile)
#         age_cls_true = Variable(age_cls_true.cuda(), volatile=volatile)
#       else:
#         inputs = Variable(inputs, volatile=volatile)
#         gender_true = Variable(gender_true, volatile=volatile)
#         age_cls_true = Variable(age_cls_true, volatile=volatile)
#
#       # 4.1.2.4 forward and get outputs
#       gender_out, age_cls_out = self.model(inputs)
#       _, gender_pred = torch.max(gender_out, 1)
#       _, age_cls_pred = torch.max(age_cls_out, 1)
#       gender_true = gender_true.view(-1)
#       age_cls_true = age_cls_true.view(-1, 99)
#
#       # 4.1.2.5 get loss
#       # print(age_cls_out.size(), age_cls_true.size(), loss_weight.size())
#       gender_loss = self.gender_criterion(gender_out, gender_true)
#       age_cls_loss = self.age_cls_criterion(age_cls_out, age_cls_true)
#       # age_rgs_loss  = self.age_rgs_criterion(age_rgs_pred, age_rgs_true)
#
#       # *Note: reduce some age loss and gender loss
#       #         enforce the model to focuse on reducing
#       #         age classification loss
#       gender_loss *= self.reduce_gen_loss
#       # age_rgs_loss *= self.reduce_age_mae
#
#       # loss = gender_loss + age_rgs_loss + age_cls_loss
#       # loss = age_cls_loss
#       loss = gender_loss + age_cls_loss
#
#       gender_loss_perc = 100 * (gender_loss / loss).cpu().data.numpy()[0]
#       age_cls_loss_perc = 100 * (age_cls_loss / loss).cpu().data.numpy()[0]
#       # age_rgs_loss_perc = 100 * (age_rgs_loss / loss).cpu().data.numpy()[0]
#
#       age_rgs_loss_perc = 0
#       weigh = np.linspace(1, 99, 99)
#       age_cls_raw = age_cls_out.cpu().data.numpy()
#       age_cls_raw = np.sum(age_cls_raw * weigh, axis=1)
#       age_rgs_true = age_rgs_true.view(-1)
#       age_rgs_true = age_rgs_true.cpu().numpy() * self.age_divide
#       age_rgs_loss = np.mean(np.abs(age_cls_raw - age_rgs_true))
#
#       # 4.1.2.7 statistics
#       gender_pred = gender_pred.cpu().data.numpy()
#       gender_true = gender_true.cpu().data.numpy()
#       batch_gender_tp = np.sum(gender_pred == gender_true)
#
#       age_cls_pred = age_cls_pred.cpu().data.numpy()
#       age_cls_true = age_rgs_true
#       batch_age_tp = np.sum(np.abs(age_cls_true - age_cls_pred) <= 2)  # if true, MAE < 5
#
#       epoch_age_mae += age_rgs_loss * inputs.size(0)
#       epoch_age_tp += batch_age_tp
#       epoch_gender_tp += batch_gender_tp
#
#       # 4.1.2.8 print info for each bach done
#       print("|| {:.2f}% {}/{} || LOSS = {:.2f} || DISTR% {:.0f} : {:.0f} : {:.0f} "
#             "|| AMAE/AACC±2/GACC = {:.2f} / {:.2f}% / {:.2f}% "
#             "|| LR {} || ETA {:.0f}s "
#             .format(100 * processed_data / self.dataset_sizes[phase],
#                     processed_data,
#                     self.dataset_sizes[phase],
#                     loss.cpu().data.numpy()[0],
#                     age_rgs_loss_perc,
#                     age_cls_loss_perc,
#                     gender_loss_perc,
#                     age_rgs_loss,
#                     # self.age_divide * age_rgs_loss.cpu().data.numpy()[0],
#                     100 * batch_age_tp / inputs.size(0),
#                     100 * batch_gender_tp / inputs.size(0),
#                     self.lr_rate,
#                     (self.dataset_sizes[phase] - processed_data) * (time.time() - epoch_start_time) / processed_data,
#             end='\r'))
#
#       # 4.1.2.9 unlink cuda variables and free up mem
#       del inputs, gender_true, age_rgs_true, age_cls_true
#       del age_cls_loss, age_rgs_loss, loss  # , gen_loss
#       del gender_loss_perc, age_cls_loss_perc, age_rgs_loss_perc
#
#     # 4.1.3.0 epoch done
#     epoch_gender_acc = epoch_gender_tp / self.dataset_sizes[phase]
#     epoch_age_acc = epoch_age_tp / self.dataset_sizes[phase]
#     epoch_age_mae /= self.dataset_sizes[phase]
#
#     # 4.1.4.0 print info after each epoch done
#     print('\n--{} Done! '
#           '|| AMAE/AACC±2/GACC = {:.2f} / {:.2f}% / {:.2f}%  '
#           '|| COST {:.0f}s'
#           .format(phase.upper(),
#                   epoch_age_mae,
#                   # self.age_divide * epoch_age_mae,
#                   100 * epoch_age_acc,
#                   100 * epoch_gender_acc,
#                   time.time() - epoch_start_time))
# """

    def getAgeGender(self,
                     img,
                     transformed=False,
                     return_all_faces=True,
                     return_info=False):
        """
    evaluation/test funtion
    :param img: str or numpy array represent the image
    :param transformed: if the image is transformed into standarlized pytorch image.
            applicable when using this in train loop
    :param return_all_faces: if set, return prediction results of all faces detected.
            set to False if it's known that all images comtain only 1 face
    :param return_info: if set, return a list of rects (x, y, w, h) represents loc of faces
    :return: a list of [gender_pred, age_pred]
    """
        # load model params
        if not self.weight_loaded:
            path = self.checkpoint_best if self.load_best else self.checkpoint_last
            checkpoint = torch.load(
                path, map_location='gpu' if self.use_gpu else 'cpu')
            self.soft_load_statedic(checkpoint['state_dic'])
            # self.model.load_state_dict(checkpoint['state_dic'])
            self.model.train(False)
            self.weight_loaded = True

        # load images if not provided
        if type(img) == str:
            img = cv2.cvtColor(cv2.imread(img), cv2.COLOR_BGR2RGB)

        # get faces and rects
        aligned = self.aligner.getAligns(img, return_info=return_info)
        if return_info:
            aligned, rects, scores = aligned
        if not len(aligned):  # no face detected
            scores = [1]
            rects = [(0, 0, img.shape[0], img.shape[1])]
            faces = [img]
        else:
            faces = aligned
        if not return_all_faces:
            faces = faces[0]
        faces = [transforms.ToPILImage()(fc) for fc in faces]
        if not transformed:
            faces = [self.transformer['val'](fc) for fc in faces]

        # get predictions of each face
        preds = self.model.evaluate(faces)

        if return_info:
            return preds, rects, scores
        return preds
Exemple #6
0
class AgePredModel:
    """ train/test class for age/gender prediction """
    def __init__(self, train_from_scratch=False, eval_use_only=False):
        """
    age prediction model, provide APIs for public uses
    :param onlyTrainLastLayers: is set, freeze params of previous layers
    :param train_from_scratch: if set, do not load pretrained params
    :param eval_use_only: if set, model will not load training/testing data
    """
        self.from_scratch = train_from_scratch
        self.model = RetrainedResnetModel()
        self.use_gpu = torch.cuda.is_available()
        self.transformer = resnet_transformer()
        self.checkpoint_best = config.model + "best_torch.nn"
        self.checkpoint_last = config.model + "last_torch.nn"
        self.latest_weights = config.model + "latest_torch.nn"
        self.csv_path = config.model + "log.csv"

        self.batch_size = int(parser['TRAIN']['batch_size'])
        self.num_epochs = int(parser['TRAIN']['num_epochs'])
        self.loading_jobs = int(parser['TRAIN']['jobs_to_load_data'])
        self.max_no_reduce = int(parser['TRAIN']['max_no_reduce'])

        self.weight_decay = float(parser['TRAIN']['weight_decay'])
        self.age_divide = float(parser['DATA']['age_divide'])
        self.min_lr_rate = float(parser['TRAIN']['min_lr_rate'])
        self.lr_reduce_by = float(parser['TRAIN']['lr_reduce_by'])
        self.lr_rate = float(parser['TRAIN']['init_lr_rate'])

        self.loaded = False
        self.age_criterion = nn.L1Loss()
        self.gender_criterion = nn.NLLLoss()
        self.aligner = FaceAligner()

        if self.use_gpu:
            self.model = self.model.cuda()
            self.age_criterion = self.age_criterion.cuda()
            self.gender_criterion = self.gender_criterion.cuda()

        columns = [
            'Timstamp', 'Epoch', 'Phase', 'Gender Acc', 'Age MAE',
            'Best Gender Acc', 'Best Age MAE', 'Lr_rate'
        ]
        self.csv_checkpoint = pd.DataFrame(data=[], columns=columns)
        if not self.from_scratch and os.path.exists(self.csv_path):
            self.csv_checkpoint = pd.read_csv(self.csv_path)

        if not eval_use_only:
            self.load_data()

    def load_data(self):
        print("[AgePredModel] load_data: start loading...")
        image_datasets = {
            x: FaceDataset(config.pics + x + '/', self.transformer[x])
            for x in ['train', 'val']
        }
        self.dataloaders = {
            x: torch.utils.data.DataLoader(image_datasets[x],
                                           batch_size=self.batch_size,
                                           shuffle=True,
                                           num_workers=self.loading_jobs)
            for x in ['train', 'val']
        }
        self.dataset_sizes = {
            x: len(image_datasets[x])
            for x in ['train', 'val']
        }
        print(
            "[AgePredModel] load_data: Done! Get {} for train and {} for test!"
            .format(self.dataset_sizes['train'], self.dataset_sizes['val']))
        print("[AgePredModel] load_data: loading finished !")

    # TODO: Double Check the Training LOOP !!!
    def train_model(self):
        print(self.model)
        print("[AgePredModel] train_model: Start training...")
        since = time.time()

        # 1.0.0.0 Define Perams
        best_model_wts = copy.deepcopy(self.model.state_dict())
        best_gender_acc, best_age_loss = .0, 100
        didnt_reduce_rounds = 0

        # 2.0.0.0 init optimizer
        self.optimizer = optim.Adam(filter(lambda p: p.requires_grad,
                                           self.model.parameters()),
                                    lr=self.lr_rate,
                                    weight_decay=self.weight_decay)

        # 3.0.0.0 load weights if possible
        if not self.from_scratch and os.path.exists(self.checkpoint_best):
            try:
                checkpoint = torch.load(self.checkpoint_best)
                self.model.load_state_dict(checkpoint['state_dic'],
                                           strict=False)
                best_gender_acc = checkpoint['best_gender_acc']
                best_age_loss = checkpoint['best_age_loss']
                best_model_wts = checkpoint['state_dic']
                self.optimizer.load_state_dict(checkpoint['optimizer'])
                print("[train_model] Load Weights Successful")
            except:
                pass

        # 4.0.0.0 start each epoch
        for epoch in range(self.num_epochs):
            print('Start Epoch {}/{} ...'.format(epoch, self.num_epochs - 1))
            print('-' * 10)

            # 4.1.0.0 loop over training and validation phase
            for phase in ['train', 'val']:
                # 4.1.1.0 shift train/eval model
                self.model.train(phase == 'train')
                torch.cuda.empty_cache()

                epoch_age_loss = 0.0
                epoch_gender_tp = 0
                processed_data = 0

                # 4.1.2.0 Iterate over data.
                epoch_start_time = time.time()
                for data in self.dataloaders[phase]:
                    # 4.1.2.1 get the inputs and labels
                    inputs, gender_true, age_true = data
                    processed_data += self.batch_size

                    # 4.1.2.2 wrap inputs&oputpus into Variable
                    #         NOTE: set voloatile = True when
                    #         doing evaluation helps reduce
                    #         gpu mem usage.
                    volatile = phase == 'val'
                    if self.use_gpu:
                        inputs = Variable(inputs.cuda(), volatile=volatile)
                        gender_true = Variable(gender_true.cuda(),
                                               volatile=volatile)
                        age_true = Variable(age_true.cuda(), volatile=volatile)
                    else:
                        inputs = Variable(inputs, volatile=volatile)
                        gender_true = Variable(gender_true, volatile=volatile)
                        age_true = Variable(age_true, volatile=volatile)

                    # 4.1.2.3 zero gradients
                    self.optimizer.zero_grad()

                    # 4.1.2.4 forward and get outputs
                    gender_out, age_pred = self.model(inputs)
                    _, gender_pred = torch.max(gender_out, 1)
                    gender_true = gender_true.view(-1)

                    # 4.1.2.5 get loss
                    gender_loss = self.gender_criterion(
                        gender_out, gender_true)
                    age_loss = self.age_criterion(age_pred, age_true)
                    loss = age_loss + gender_loss

                    # 4.1.2.6 backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        self.optimizer.step()

                    # 4.1.2.7 statistics
                    gender_pred = gender_pred.cpu().data.numpy()
                    gender_true = gender_true.cpu().data.numpy()
                    this_epoch_gender_tp = np.sum(gender_pred == gender_true)
                    epoch_age_loss += age_loss.data[0] * inputs.size(0)
                    epoch_gender_tp += this_epoch_gender_tp

                    # 4.1.2.8 print info for each bach
                    print(
                        "|| {:.2f}% {}/{} || Gender Loss {:.4f} || Age MAE {:.2f} || Acc {:.2f}% |"
                        "| LR {} || ETA {:.0f}s || BEST AGE {:.2f} || BEST GENDER {:.2f}% ||"
                        .format(
                            100 * processed_data / self.dataset_sizes[phase],
                            processed_data, self.dataset_sizes[phase],
                            gender_loss.cpu().data.numpy()[0],
                            self.age_divide * age_loss.cpu().data.numpy()[0],
                            100 * this_epoch_gender_tp / self.batch_size,
                            self.lr_rate,
                            (self.dataset_sizes[phase] - processed_data) *
                            (time.time() - epoch_start_time) / processed_data,
                            self.age_divide * best_age_loss,
                            100 * best_gender_acc))

                    # 4.1.2.9 free up mem
                    del inputs, gender_true, age_true
                    del gender_loss, age_loss, loss

                # 4.1.3.0 epoch done
                epoch_gender_acc = epoch_gender_tp / self.dataset_sizes[phase]
                epoch_age_loss /= self.dataset_sizes[phase]

                # 4.1.4.0 print info after epoch done
                print(
                    '\n\n{} {}/{} Done! \t\t Age MAE = {:.2f} \t\t Gender Acc {:.2f}% Lr = {:.5f}\n\n'
                    .format(phase.upper(), epoch, self.num_epochs,
                            self.age_divide * epoch_age_loss,
                            100 * epoch_gender_acc, self.lr_rate))

                # 4.1.5.0, save model weights
                if phase == 'val' and epoch_age_loss < best_age_loss:
                    best_gender_acc = epoch_gender_acc
                    best_age_loss = epoch_age_loss
                    best_model_wts = copy.deepcopy(self.model.state_dict())
                    torch.save(
                        {
                            'epoch': epoch,
                            'state_dic': best_model_wts,
                            "best_gender_acc": best_gender_acc,
                            "best_age_loss": best_age_loss,
                            "optimizer": self.optimizer.state_dict()
                        }, self.checkpoint_best)
                    didnt_reduce_rounds = 0
                    print(
                        "\n\tNew BEST FOUND!! Age Loss = {:.2f}, Gender Acc = {:.2f}\n"
                        .format(best_age_loss, best_gender_acc))
                elif phase == 'train':
                    didnt_reduce_rounds += 1
                    torch.save(
                        {
                            'epoch': epoch,
                            'state_dic': self.model.state_dict(),
                            "best_gender_acc": best_gender_acc,
                            "best_age_loss": best_age_loss,
                            "optimizer": self.optimizer.state_dict()
                        }, self.checkpoint_last)

                # 4.1.6.0 save csv logging file
                self.csv_checkpoint.loc[len(self.csv_checkpoint)] = [
                    str(datetime.datetime.now()), epoch, phase,
                    epoch_gender_acc, self.age_divide * epoch_age_loss,
                    best_gender_acc, self.age_divide * best_age_loss,
                    self.lr_rate
                ]
                self.csv_checkpoint.to_csv(config.model + "log.csv",
                                           index=False)

                # 4.1.7.0 reduce learning rate if nessessary
                if phase == "train" \
                        and didnt_reduce_rounds >= self.max_no_reduce\
                        and self.lr_rate > self.min_lr_rate:
                    self.lr_rate /= self.lr_reduce_by
                    print(
                        "\n[didnt_reduce_rounds] Reduce Learning Rate From {} --> {} \n"
                        .format(self.lr_rate * self.lr_reduce_by,
                                self.lr_rate))
                    for param_group in self.optimizer.param_groups:
                        param_group['lr'] = self.lr_rate
                    didnt_reduce_rounds = 0

            # 4.2.0.0 train/val loop ends

        # 5.0.0.0 Trainning Complete!
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(
            time_elapsed // 60, time_elapsed % 60))
        print('Best val enderACC: {:.4f} AgeMAE: {:.0f}'.format(
            best_gender_acc, torch.sqrt(best_age_loss)))

        # 6.0.0.0 load best model weights
        self.model.load_state_dict(best_model_wts)
        return self.model

    def getAgeGender(self,
                     img,
                     transformed=False,
                     return_all_faces=True,
                     return_info=False):
        """
    evaluation/test funtion
    :param img: str or numpy array represent the image
    :param transformed: if the image is transformed into standarlized pytorch image
            applicable when using this in train loop
    :param return_all_faces: if set, return prediction results of all faces detected,
            applicable if it's known that all images comtain only 1 face
    :param return_info: if set, return a list of [ (x, y, w, h) ] represents loc of faces
    :return: a list of [gender_pred, age_pred]
    """
        # load model params
        if not self.loaded:
            checkpoint = torch.load(
                self.checkpoint_best,
                map_location='gpu' if self.use_gpu else 'cpu')
            self.model.load_state_dict(checkpoint['state_dic'])
            self.model.train(False)
            self.loaded = True

        # load images if not provided
        if type(img) == str:
            img = cv2.imread(img)
        # img = deepcopy(img)

        aligned = self.aligner.getAligns(img, return_info=return_info)

        # TODO: Check this default rects
        if return_info:
            aligned, rects, scores = aligned

        if not len(aligned):
            scores = [1]
            rects = [(0, 0, img.shape[0], img.shape[1])]
            faces = [img]
        else:
            faces = aligned

        if not return_all_faces:
            faces = faces[0]

        faces = [transforms.ToPILImage()(fc) for fc in faces]

        if not transformed:
            faces = [self.transformer['val'](fc) for fc in faces]

        preds = []
        for face in faces:
            face = Variable(torch.unsqueeze(face, 0))
            gender_out, age_pred = self.model(face)
            _, gender_pred = torch.max(gender_out, 1)
            gender_pred = gender_pred.cpu().data.numpy()[0]
            age_pred = 10 * age_pred.cpu().data.numpy()[0][0]
            preds += [(gender_pred, age_pred)]

        if return_info:
            return preds, rects, scores
        return preds