FCHD/train.py at master · embracesource-cv-com/FCHD · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# -*- coding: utf-8 -*-
from torch.utils.data import DataLoader
from torchvision import transforms
from data.dataset import HeadDataset
from data.preprocess import Rescale, Normalize, inverse_normalize
from config import cfg
from utils.visualize import visdom_bbox
import numpy as np
import os
from networks.detector import HeadDetector
from trainer import Trainer
from utils import tools
import time


def train():
    # Load data
    train_annots_path = os.path.join(cfg.DATASET_DIR, cfg.TRAIN_ANNOTS_FILE)
    val_annots_path = os.path.join(cfg.DATASET_DIR, cfg.VAL_ANNOTS_FILE)
    transform = transforms.Compose([Rescale(), Normalize()])
    train_dataset = HeadDataset(cfg.DATASET_DIR, train_annots_path, transform)
    val_dataset = HeadDataset(cfg.DATASET_DIR, val_annots_path, transform)

    print('[INFO] Load datasets.\n Training set size:{}, Verification set size:{}'
          .format(len(train_dataset), len(val_dataset)))

    train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True)
    val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=True)

    # Create HeadDetector instance and Trainer instance
    head_detector = HeadDetector(ratios=cfg.ANCHOR_RATIOS, scales=cfg.ANCHOR_SCALES)
    trainer = Trainer(head_detector).cuda()

    print('[INFO] Start training...')
    for epoch in range(cfg.EPOCHS):
        trainer.reset_meters()
        for i, data in enumerate(train_dataloader, 1):
            img, boxes, scale = data['img'], data['boxes'], data['scale']
            img = img.cuda().float()
            scale = scale.item()

            # Forward pass and backward pass
            trainer.train_step(img, boxes, scale)

            # Visualize on visdom
            if i % cfg.PLOT_INTERVAL == 0:
                trainer.vis.plot_many(trainer.get_meter_data())
                origin_img = inverse_normalize(img[0].cpu().numpy())
                gt_img = visdom_bbox(origin_img, boxes[0].cpu().numpy())
                trainer.vis.img('gt_img', gt_img)
                preds, _ = head_detector(img, scale)
                pred_img = visdom_bbox(origin_img, preds)
                trainer.vis.img('pred_img', pred_img)
                trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm')

        # Evaluation
        avg_accuracy = evaluate(val_dataloader, head_detector)
        print("[INFO] Epoch {} of {}.".format(epoch + 1, cfg.EPOCHS))
        print("\tValidate average accuracy: {:.3f}".format(avg_accuracy))

        # Save current model
        time_str = time.strftime('%m%d%H%M')
        save_path = os.path.join(cfg.MODEL_DIR, 'checkpoint_{}_{:.3f}.pth'.format(time_str, avg_accuracy))
        trainer.save(save_path)

        # Learning rate decay
        if epoch == 8:
            trainer.scale_lr()


def evaluate(val_dataloader, head_detector):
    img_counts = 0
    accuracy = 0.0

    for data in val_dataloader:
        img, boxes, scale = data['img'], data['boxes'], data['scale']
        img, boxes = img.cuda().float(), boxes.cuda()
        scale = scale.item()

        preds, _ = head_detector(img, scale)
        gts = boxes[0].cpu().numpy()
        if len(preds) == 0:
            img_counts += 1
        else:
            ious = tools.calc_ious(preds, gts)
            max_ious = ious.max(axis=1)
            correct_counts = len(np.where(max_ious >= 0.5)[0])
            gt_counts = len(gts)
            accuracy += correct_counts / gt_counts
            img_counts += 1

    avg_accuracy = accuracy / img_counts
    return avg_accuracy


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    train()