def main(config):
# 选择使用模式
if config.mode == 'train':
# 获取DataLoader
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中保存文件夹名称
config.save_folder = "%s/run-%d" % (config.save_folder, run)
# 开始进入training状态
train = Solver(train_loader, None, config)
train.train()
elif config.mode == 'test':
# 选择数据集路径
config.test_root, config.test_list = get_test_info(config.sal_mode)
# 获取DataLoader
test_loader = get_loader(config, mode='test')
# 如果文件路径不存在
if not os.path.exists(config.test_fold):
os.mkdir(config.test_fold)
# 开始进入testing状态
test = Solver(None, test_loader, config)
test.test(test_mode=config.test_mode)
else:
# 非法输入
raise IOError("illegal input!!!")
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(test_mode=config.test_mode)
else:
raise IOError("illegal input!!!")