def train(cfg):
model = build_model(cfg)
# model.fix_bn()
model = nn.DataParallel(model)
torch.backends.cudnn.benchmark = True
optimizer = make_optimizer(cfg, model)
criterion = make_criterion(cfg)
scheduler = make_lr_scheduler(cfg, optimizer)
metrics = make_metrics(cfg)
train_loader = make_dataloader(cfg, is_train=True)
val_loader = make_dataloader(cfg, is_train=False)
cfg.TOOLS.image_n = 3
#image_3_dataloader = make_inference_dataloader(cfg=cfg)
image_3_dataloader = None
cfg.TOOLS.image_n = 4
#image_4_dataloader = make_inference_dataloader(cfg=cfg)
image_4_dataloader = None
do_train(cfg,
model=model,
train_loader=train_loader,
val_loader=val_loader,
optimizer=optimizer,
scheduler=scheduler,
loss_fn=criterion,
metrics=metrics,
image_3_dataloader=image_3_dataloader,
image_4_dataloader=image_4_dataloader)
free_device_ids = get_free_device_ids()
max_num_devices = cfg['max_num_devices']
if len(free_device_ids) >= max_num_devices:
free_device_ids = free_device_ids[:max_num_devices]
# print(free_device_ids)
master_device = free_device_ids[0]
model = nn.DataParallel(model,
device_ids=free_device_ids).cuda(master_device)
# save_dir
curr_time = datetime.datetime.now()
time_str = datetime.datetime.strftime(curr_time, '%Y%m%d_')
model_tag = (os.path.split(arg['load_path'])[1]).split('.')[0]
save_dir = os.path.join(r'../exp/', time_str + model_tag)
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
print("Save_dir :", save_dir)
dataloader = make_dataloader(cfg['query_pipeline'])
dict_query = ADD_TTA_inference_extract_features(cfg,
model,
dataloader,
device_ids=free_device_ids)
dict_to_file(os.path.join(save_dir, "query.json"), dict_query)
dataloader = make_dataloader(cfg['gallery_pipeline'])
dict_gallery = ADD_TTA_inference_extract_features(
cfg, model, dataloader, device_ids=free_device_ids)
dict_to_file(os.path.join(save_dir, "gallery.json"), dict_gallery)
if __name__ == "__main__":
# 若更新了load_arg函数，需要对应更新merage_from_arg()
arg = vars(load_arg())
# 待修改
config_file = arg["CONFIG_FILE"]
config_file = config_file.replace("../","").replace(".py","").replace('/','.')
exec(r"from {} import config as cfg".format(config_file))
# if arg['MODEL.LOAD_PATH'] != None: #优先级：arg传入命令 >model中存的cfg > config_file
# cfg = torch.load(arg['MODEL.LOAD_PATH'])['cfg']
cfg = merage_from_arg(cfg,arg)
print(cfg)
cfg_copy = copy.deepcopy(cfg)
train_dataloader = make_dataloader(cfg['train_pipeline'])
curr_time = datetime.datetime.now()
time_str = datetime.datetime.strftime(curr_time,'%Y%m%d_')
save_dir = os.path.join(cfg['save_dir'],time_str+cfg['tag'])
log_dir = os.path.join(cfg['log_dir'],"log_"+time_str+cfg['tag'])
cfg['save_dir'] = save_dir
cfg['log_dir'] = log_dir
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
if not os.path.isdir(log_dir):
os.makedirs(log_dir)
print("Save_dir :",save_dir)
print("Log_dir :", log_dir)
from PIL import ImageFile
from PIL import Image
from torch.utils.data import DataLoader
from models.model import Generator, Discriminator
from utils.utils import calc_optical_flow, weights_init, make_G_input, make_D_input, save_video
from options.train_options import TrainOptions
from data.dataloader import make_dataloader
ImageFile.LOAD_TRUNCATED_IMAGES = True
opt = TrainOptions().parse()
localtime = time.asctime(time.localtime(time.time()))
start_time = time.time()
train_loader, valid_loader = make_dataloader(opt)
# Decide which device we want to run on
device = torch.device("cuda:0" if (
torch.cuda.is_available() and opt.ngpu > 0) else "cpu")
# Create the generator
netG = Generator().to(device)
if (device.type == 'cuda') and (opt.ngpu > 1):
netG = nn.DataParallel(netG, list(range(opt.ngpu)))
netG.apply(weights_init)
print(netG)
# Create the Discriminator
netD = Discriminator().to(device)
if (device.type == 'cuda') and (opt.ngpu > 1):
# print(free_device_ids)
master_device = free_device_ids[0]
model = nn.DataParallel(model,
device_ids=free_device_ids).cuda(master_device)
# save_dir
curr_time = datetime.datetime.now()
time_str = datetime.datetime.strftime(curr_time, '%Y%m%d_')
model_tag = (os.path.split(arg['load_path'])[1]).split('.')[0]
save_dir = os.path.join(r'../exp/', time_str + model_tag)
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
print("Save_dir :", save_dir)
image_3_pipeline = make_dataloader(cfg['test_image3_pipeline'])
image_3_predict = tta_forward(
image_3_pipeline,
model,
device=master_device,
png_shape=cfg['test_image3_pipeline']['image_shape'])
pil_image = Image.fromarray(image_3_predict)
pil_image.save(os.path.join(save_dir, "image_" + str(3) + "_predict.png"))
del image_3_pipeline, image_3_predict
image_4_pipeline = make_dataloader(cfg['test_image4_pipeline'])
image_4_predict = tta_forward(
image_4_pipeline,
model,
device=master_device,
png_shape=cfg['test_image4_pipeline']['image_shape'])
if __name__ == "__main__":
# 若更新了load_arg函数，需要对应更新merage_from_arg()
arg = vars(load_arg())
if arg['MODEL.LOAD_PATH'] != None: #优先级：arg传入命令 >model中存的cfg > config_file
cfg = torch.load(arg['MODEL.LOAD_PATH'])['cfg']
# 待修改
config_file = arg["CONFIG_FILE"]
config_file = config_file.replace("../", "").replace(".py",
"").replace('/', '.')
exec(r"from {} import config as cfg".format(config_file))
cfg = merage_from_arg(cfg, arg)
model = bulid_model(cfg['model'], cfg['pretrain'])
# 默认开启多GPU
if cfg['multi_gpu']:
device_ids = cfg['device_ids']
model = nn.DataParallel(model, device_ids=device_ids)
curr_time = datetime.datetime.now()
time_str = datetime.datetime.strftime(curr_time, '%Y%m%d_')
model_tag = (os.path.split(arg['MODEL.LOAD_PATH'])[1]).split('.')[0]
save_dir = os.path.join(r'../exp/', time_str + model_tag)
if not os.path.isdir(save_dir):
os.makedirs(save_dir)
print("Save_dir :", save_dir)
cfg['save_dir'] = save_dir
dataloader = make_dataloader(cfg['test_pipeline'])
inference_forward(cfg, model, dataloader, tta_flag=True)
zipDir(save_dir, save_dir + ".zip") # 压缩至根目录