-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
87 lines (72 loc) · 5.08 KB
/
main.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
""" Main function for this project. """
import os
import argparse
import numpy as np
from trainer.trainer import Trainer
from utils.gpu_tools import occupy_memory
if __name__ == '__main__':
parser = argparse.ArgumentParser()
### Basic parameters
parser.add_argument('--gpu', default='0', help='the index of GPU')
parser.add_argument('--dataset', default='cifar100', type=str, choices=['cifar100', 'imagenet_sub', 'imagenet'])
parser.add_argument('--data_dir', default='./data', type=str)
parser.add_argument('--baseline', default='lucir', type=str, choices=['lucir', 'icarl'], help='baseline method')
parser.add_argument('--ckpt_label', type=str, default='exp01', help='the label for the checkpoints')
parser.add_argument('--ckpt_dir_fg', type=str, default='-', help='the checkpoint file for the 0-th phase')
parser.add_argument('--resume_fg', action='store_true', help='resume 0-th phase model from the checkpoint')
parser.add_argument('--resume', action='store_true', help='resume from the checkpoints')
parser.add_argument('--num_workers', default=1, type=int, help='the number of workers for loading data')
parser.add_argument('--random_seed', default=1993, type=int, help='random seed')
parser.add_argument('--train_batch_size', default=128, type=int, help='the batch size for train loader')
parser.add_argument('--test_batch_size', default=100, type=int, help='the batch size for test loader')
parser.add_argument('--eval_batch_size', default=128, type=int, help='the batch size for validation loader')
parser.add_argument('--disable_gpu_occupancy', action='store_false', help='disable GPU occupancy')
### Network architecture parameters
parser.add_argument('--branch_mode', default='dual', type=str, choices=['dual', 'single'], help='the branch mode for AANets')
parser.add_argument('--branch_1', default='ss', type=str, choices=['ss', 'fixed', 'free'], help='the network type for the first branch')
parser.add_argument('--branch_2', default='free', type=str, choices=['ss', 'fixed', 'free'], help='the network type for the second branch')
parser.add_argument('--imgnet_backbone', default='resnet18', type=str, choices=['resnet18', 'resnet34'], help='network backbone for ImageNet')
### Incremental learning parameters
parser.add_argument('--num_classes', default=100, type=int, help='the total number of classes')
parser.add_argument('--nb_cl_fg', default=50, type=int, help='the number of classes in the 0-th phase')
parser.add_argument('--nb_cl', default=10, type=int, help='the number of classes for each phase')
parser.add_argument('--nb_protos', default=20, type=int, help='the number of exemplars for each class')
parser.add_argument('--epochs', default=160, type=int, help='the number of epochs')
parser.add_argument('--dynamic_budget', action='store_true', help='using dynamic budget setting')
parser.add_argument('--fusion_lr', default=1e-8, type=float, help='the learning rate for the aggregation weights')
### General learning parameters
parser.add_argument('--lr_factor', default=0.1, type=float, help='learning rate decay factor')
parser.add_argument('--custom_weight_decay', default=5e-4, type=float, help='weight decay parameter for the optimizer')
parser.add_argument('--custom_momentum', default=0.9, type=float, help='momentum parameter for the optimizer')
parser.add_argument('--base_lr1', default=0.1, type=float, help='learning rate for the 0-th phase')
parser.add_argument('--base_lr2', default=0.1, type=float, help='learning rate for the following phases')
### LUCIR parameters
parser.add_argument('--the_lambda', default=5, type=float, help='lamda for LF')
parser.add_argument('--dist', default=0.5, type=float, help='dist for margin ranking losses')
parser.add_argument('--K', default=2, type=int, help='K for margin ranking losses')
parser.add_argument('--lw_mr', default=1, type=float, help='loss weight for margin ranking losses')
### iCaRL parameters
parser.add_argument('--icarl_beta', default=0.25, type=float, help='beta for iCaRL')
parser.add_argument('--icarl_T', default=2, type=int, help='T for iCaRL')
### CSCCT specific parameters
parser.add_argument('--csc', action='store_true', default=False)
parser.add_argument('--csc_weight', type=float, default=3)
parser.add_argument('--ct', action='store_true', default=False)
parser.add_argument('--ct_weight', type=float, default=1.5)
parser.add_argument('--ct_temperature', type=float, default=2)
the_args = parser.parse_args()
# Checke the number of classes, ensure they are reasonable
assert(the_args.nb_cl_fg % the_args.nb_cl == 0)
assert(the_args.nb_cl_fg >= the_args.nb_cl)
# Print the parameters
print(the_args)
# Set GPU index
os.environ['CUDA_VISIBLE_DEVICES'] = the_args.gpu
print('Using gpu:', the_args.gpu)
# Occupy GPU memory in advance
if the_args.disable_gpu_occupancy:
occupy_memory(the_args.gpu)
print('Occupy GPU memory in advance.')
# Set the trainer and start training
trainer = Trainer(the_args)
trainer.train()