-
Notifications
You must be signed in to change notification settings - Fork 1
/
mgtn.py
134 lines (117 loc) · 6.66 KB
/
mgtn.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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import argparse
from engine import *
from models import *
from coco import *
from util import *
import neptune
parser = argparse.ArgumentParser(description='MGTN Training')
parser.add_argument('data', metavar='DIR',
help='path to dataset (e.g. data/')
parser.add_argument('--image-size', '-i', default=448, type=int,
metavar='N', help='image size (default: 224)')
parser.add_argument('-j', '--workers', default=8, type=int, metavar='N',
help='number of data loading workers (default: 8)')
parser.add_argument('--epochs', default=200, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--epoch_step', default=[30], type=int, nargs='+',
help='number of epochs to change learning rate')
parser.add_argument('--device_ids', default=[0], type=int, nargs='+',
help='number of epochs to change learning rate')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=16, type=int,
metavar='N', help='mini-batch size (default: 16)')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
metavar='LR', help='initial learning rate')
parser.add_argument('--lrd', '--learning-rate-decay', default=0.1, type=float,
metavar='LRD', help='learning rate decay')
parser.add_argument('--lrp', '--learning-rate-pretrained', default=0.1, type=float,
metavar='LR', help='learning rate for pre-trained layers')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=0, type=int,
metavar='N', help='print frequency (default: 10)')
parser.add_argument('--embedding', default='model/embedding/coco_glove_word2vec_80x300.pkl',
type=str, metavar='EMB', help='path to embedding (default: glove)')
parser.add_argument('--embedding-length', default=300, type=int, metavar='EMB',
help='embedding length (default: 300)')
parser.add_argument('--adj-file', default='model/topology/coco_adj.pkl', type=str, metavar='ADJ',
help='Adj file (default: model/topology/coco_adj.pkl')
parser.add_argument('--adj-strong-threshold', default=0.4, type=float, metavar='ADJTS',
help='Adj strong threshold (default: 0.4)')
parser.add_argument('--adj-weak-threshold', default=0.2, type=float, metavar='ADJTW',
help='Adj weak threshold (default: 0.2)')
parser.add_argument('--mod-file', default='model/topology/coco_mod.txt', type=str, metavar='MOD',
help='Adj file (default: model/topology/coco_mod.txt')
parser.add_argument('--mlt', '--multi-learning-threshold', default=0.999, type=float, metavar='MLT',
help='Multi-learning threshold (default: 0.999)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
help='evaluate model on validation set')
parser.add_argument('-n', '--neptune', dest='neptune', action='store_true',
help='run with neptune')
parser.add_argument('--neptune-path', default='neptune.txt', type=str, metavar='PATH',
help='Neptune API keys (default: neptune.txt)')
parser.add_argument('--exp-name', dest='exp_name', default='coco', type=str, metavar='COCO2014',
help='Name of experiment to have different location to save checkpoints')
def main_coco():
global args, best_prec1, use_gpu
args = parser.parse_args()
use_gpu = torch.cuda.is_available()
torch.backends.cudnn.benchmark = False
if args.neptune:
with open(args.neptune_path, 'r') as f:
nep = f.readlines()
neptune.init(nep[0].strip(), api_token=nep[1].strip())
neptune.create_experiment(params=vars(args), upload_source_files=[
'demo_coco_gcn.py', 'coco.py', 'engine.py', 'models.py', 'util.py'])
train_dataset = COCO2014(args.data, phase='train', inp_name=args.embedding)
val_dataset = COCO2014(args.data, phase='val', inp_name=args.embedding)
num_classes = 80
print('Embedding:', args.embedding, '(', args.embedding_length, ')')
print('Adjacency file:', args.adj_file)
print('Adjacency Strong Threshold:', args.adj_strong_threshold)
print('Adjacency Weak Threshold:', args.adj_weak_threshold)
print('Modularity file:', args.mod_file)
if args.adj_strong_threshold < args.adj_weak_threshold:
args.adj_weak_threshold = args.adj_strong_threshold
model = mgtn_resnet(num_classes=num_classes,
t1=args.adj_strong_threshold,
t2=args.adj_weak_threshold,
adj_file=args.adj_file,
mod_file=args.mod_file,
in_channel=args.embedding_length,
ml_threshold=args.mlt)
# define loss function (criterion)
criterion = nn.MultiLabelSoftMarginLoss()
# define optimizer
optimizer = torch.optim.SGD(model.get_config_optim(args.lr, args.lrp),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
model_path = "checkpoint/coco/%s" % args.exp_name
if not os.path.exists(model_path):
os.makedirs(model_path)
state = {'batch_size': args.batch_size,
'image_size': args.image_size,
'max_epochs': args.epochs,
'evaluate': args.evaluate,
'resume': args.resume,
'num_classes': num_classes,
'difficult_examples': True,
'save_model_path': model_path,
'workers': args.workers,
'epoch_step': args.epoch_step,
'lr': args.lr,
'lr_decay': args.lrd,
'mlt': args.mlt,
'device_ids': args.device_ids,
'evaluate': args.evaluate,
'neptune': args.neptune}
engine = GCNMultiLabelMAPEngine(state)
engine.learning(model, criterion, train_dataset, val_dataset, optimizer)
if __name__ == '__main__':
main_coco()