-
Notifications
You must be signed in to change notification settings - Fork 4
/
main_joint.py
102 lines (86 loc) · 4.49 KB
/
main_joint.py
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
99
100
101
102
"""
Code to train the Joint Object and Part Detector using Attention-based feature fusion
"""
import argparse
import numpy as np
import os
import torch
from datasets import load_data
from models import JointDetector
from train_joint import train_one_epoch, evaluate
from utils import set_all_seeds
########## CMDLINE ARGS ##########
parser = argparse.ArgumentParser('Train Model')
parser.add_argument('-d', '--device', type=int, default=0)
parser.add_argument('-dir', '--data_dir', type=str, default='data/VOCdevkit/VOC2010/')
parser.add_argument('-tr', '--train_split', type=str, default='animals_train')
parser.add_argument('-val', '--val_split', type=str, default='animals_val')
parser.add_argument('-ocf', '--obj_class2ind_file', type=str, default='animals_object_class2ind')
parser.add_argument('-pcf', '--part_class2ind_file', type=str, default='animals_part_mergedclass2ind')
parser.add_argument('-e', '--n_epochs', type=int, default=30)
parser.add_argument('-lr', '--learning_rate', type=float, default=1e-3)
parser.add_argument('-bs', '--batch_size', type=int, default=1)
parser.add_argument('-wd', '--weight_decay', type=float, default=1e-6)
parser.add_argument('-ft', '--fusion_thresh', type=float, default=0.9)
parser.add_argument('--no_attention', dest='no_attention', action='store_true')
parser.add_argument('-nw', '--num_workers', type=int, default=0)
parser.add_argument('-ms', '--max_samples', type=int, default=-1)
args = parser.parse_args()
device = torch.device('cuda:%d' % args.device if torch.cuda.is_available() else 'cpu')
set_all_seeds(123)
########## Parameters ##########
DATA_DIR = args.data_dir
TRAIN_SPLIT = args.train_split
VAL_SPLIT = args.val_split
OBJ_CLASS2IND_FILE = args.obj_class2ind_file
PART_CLASS2IND_FILE = args.part_class2ind_file
N_EPOCHS = args.n_epochs
NUM_WORKERS = args.num_workers
MAX_SAMPLES = args.max_samples if args.max_samples > 0 else None
USE_OBJECTS = True
USE_PARTS = True
RETURN_SEPARATE_TARGETS = True
########## Hyperparameters ##########
LEARNING_RATE = args.learning_rate
BATCH_SIZE = args.batch_size
WEIGHT_DECAY = args.weight_decay
FUSION_THRESH = args.fusion_thresh
USE_ATTENTION = not args.no_attention
model_save_path = 'saved_model_joint_ft-%.2f_%s.pth' % (FUSION_THRESH, TRAIN_SPLIT)
########## Data Loaders ##########
train_loader, val_loader, obj_class2ind, obj_n_classes, part_class2ind, part_n_classes = load_data(DATA_DIR, BATCH_SIZE, TRAIN_SPLIT, VAL_SPLIT,
OBJ_CLASS2IND_FILE, USE_OBJECTS, USE_PARTS, RETURN_SEPARATE_TARGETS, PART_CLASS2IND_FILE, NUM_WORKERS, MAX_SAMPLES)
########## Create Model ##########
model = JointDetector(obj_n_classes, part_n_classes, FUSION_THRESH, USE_ATTENTION).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=LEARNING_RATE, momentum=0.9, weight_decay=WEIGHT_DECAY)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
########## Restore Saved Model If Exists ##########
start_epoch = 0
best_val_mAP = {'OBJECT': 0., 'PART': 0.}
if os.path.exists(model_save_path):
print('Restoring trained model from %s' % model_save_path)
checkpoint = torch.load(model_save_path, map_location=device)
model.load_state_dict(checkpoint['model'])
if 'lr_scheduler' in checkpoint:
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
if 'optimizer' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
best_val_mAP = checkpoint['mAP']
start_epoch = checkpoint['epoch']+1
N_EPOCHS = start_epoch + N_EPOCHS
########## Train Model ##########
for epoch in range(start_epoch, N_EPOCHS):
train_one_epoch(model, optimizer, train_loader, device, epoch, print_freq=250)
_, _, obj_det_stats, part_det_stats = evaluate(model, val_loader, device=device, print_freq=500, header='Val:')
if epoch < 10: # stop decaying after 10 epochs
lr_scheduler.step()
val_obj_det_mAP = obj_det_stats['bbox'][1] # AP @ IoU=0.5
val_part_det_mAP = part_det_stats['bbox'][1] # AP @ IoU=0.5
if val_obj_det_mAP > best_val_mAP['OBJECT']: # Save model which performs best for object detection
best_val_mAP = {'OBJECT': val_obj_det_mAP, 'PART': val_part_det_mAP}
checkpoint = {'model': model.state_dict(), #'optimizer' : optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch, 'mAP': best_val_mAP}
torch.save(checkpoint, model_save_path)
print('-'*100)
print('Best val_mAP OBJECT: %.2f%%, PART: %.2f%%' % (100*best_val_mAP['OBJECT'], 100*best_val_mAP['PART']))