-
Notifications
You must be signed in to change notification settings - Fork 4
/
main_clm.py
95 lines (78 loc) · 3.7 KB
/
main_clm.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
import argparse
import os
from dataset.dataset_saliency_augment_nofuse import get_loader
from Solver_clm import *
def get_test_info(sal_mode='e'):
if sal_mode == 'e':
image_root = './data/ECSSD/Imgs/'
image_source = './data/ECSSD/test.lst'
elif sal_mode == 'p':
image_root = './data/PASCALS/Imgs/'
image_source = './data/PASCALS/test.lst'
elif sal_mode == 'd':
image_root = './data/DUTOMRON/Imgs/'
image_source = './data/DUTOMRON/test.lst'
elif sal_mode == 'h':
image_root = './data/HKU-IS/Imgs/'
image_source = './data/HKU-IS/test.lst'
elif sal_mode == 't':
image_root = './data/DUTS-TE/Imgs/'
image_source = './data/DUTS-TE/test.lst'
return image_root, image_source
def main(config):
if config.mode == 'train':
train_loader = get_loader(config)
run = 0
while os.path.exists("%s/run-%d" % (config.save_folder, run)):
run += 1
os.mkdir("%s/run-%d" % (config.save_folder, run))
os.mkdir("%s/run-%d/models" % (config.save_folder, run))
config.save_folder = "%s/run-%d" % (config.save_folder, run)
train = Solver(train_loader, None, config)
train.train()
elif config.mode == 'test':
config.test_root, config.test_list = get_test_info(config.sal_mode)
test_loader = get_loader(config, mode='test')
if not os.path.exists(config.test_fold): os.mkdir(config.test_fold)
test = Solver(None, test_loader, config)
test.test()
else:
raise IOError("illegal input!!!")
if __name__ == '__main__':
torch.cuda.set_device(0)
vgg_path = './dataset/pretrained/vgg16_20M.pth'
resnet_path = './dataset/pretrained/resnet50_caffe.pth'
parser = argparse.ArgumentParser()
# Hyper-parameters
parser.add_argument('--n_color', type=int, default=3)
parser.add_argument('--lr', type=float, default=5e-5) # Learning rate resnet:5e-5, vgg:1e-4
parser.add_argument('--wd', type=float, default=0.0005) # Weight decay
parser.add_argument('--no-cuda', dest='cuda', action='store_false')
# Training settings
parser.add_argument('--arch', type=str, default='resnet') # resnet or vgg
parser.add_argument('--pretrained_model', type=str, default=resnet_path)
parser.add_argument('--epoch', type=int, default=24)
parser.add_argument('--batch_size', type=int, default=1) # only support 1 now
parser.add_argument('--num_thread', type=int, default=1)
parser.add_argument('--load', type=str, default='')
parser.add_argument('--save_folder', type=str, default='./results/clm')
parser.add_argument('--epoch_save', type=int, default=3)
parser.add_argument('--iter_size', type=int, default=10)
parser.add_argument('--show_every', type=int, default=1000)
# Train data
parser.add_argument('--train_root', type=str, default='./data/DUTS/DUTS-TR')
parser.add_argument('--train_list', type=str, default='./data/DUTS/DUTS-TR/train_pair.lst')
# Testing settings
parser.add_argument('--model', type=str, default=None) # Snapshot
parser.add_argument('--test_fold', type=str, default=None) # Test results saving folder
parser.add_argument('--sal_mode', type=str, default='e') # Test image dataset
# Misc
parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
config = parser.parse_args()
if not os.path.exists(config.save_folder):
os.mkdir(config.save_folder)
# Get test set info
test_root, test_list = get_test_info(config.sal_mode)
config.test_root = test_root
config.test_list = test_list
main(config)