349 lines
13 KiB
Python
349 lines
13 KiB
Python
#!/usr/bin/python
|
||
# -*- encoding: utf-8 -*-
|
||
import os
|
||
#os.environ['CUDA_VISIBLE_DEVICES'] = '1'
|
||
from logger import setup_logger
|
||
from models.model_stages import BiSeNet
|
||
# from heliushuju import Heliushuju
|
||
from heliushuju_process import Heliushuju
|
||
from loss.loss import OhemCELoss
|
||
from loss.detail_loss import DetailAggregateLoss
|
||
# from evaluation import MscEvalV0
|
||
from evaluation_process import MscEvalV0
|
||
from optimizer_loss import Optimizer
|
||
import sys
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.data import DataLoader
|
||
import torch.nn.functional as F
|
||
import torch.distributed as dist
|
||
|
||
|
||
import os.path as osp
|
||
import logging
|
||
import time
|
||
import datetime
|
||
import argparse
|
||
import json
|
||
logger = logging.getLogger()
|
||
|
||
def str2bool(v):
|
||
if v.lower() in ('yes', 'true', 't', 'y', '1'):
|
||
return True
|
||
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
|
||
return False
|
||
else:
|
||
raise argparse.ArgumentTypeError('Unsupported value encountered.')
|
||
|
||
|
||
def parse_args():
|
||
parse = argparse.ArgumentParser()
|
||
parse.add_argument('--parJson',dest = 'parJson',type = str,default = './data/Car',)
|
||
parse.add_argument('--gpuId',dest = 'gpuId',type = str,default = '0',)
|
||
parse.add_argument('--local_rank',dest = 'local_rank',type = int,default = -1 )
|
||
parse.add_argument('--n_workers_train',dest = 'n_workers_train',type = int,default = 8)
|
||
parse.add_argument('--n_workers_val',dest = 'n_workers_val',type = int,default = 2)
|
||
parse.add_argument('--n_img_per_gpu',dest = 'n_img_per_gpu',type = int,default = 8,)
|
||
parse.add_argument('--max_iter',dest = 'max_iter',type = int,default = 43000)
|
||
parse.add_argument('--save_iter_sep',dest = 'save_iter_sep',type = int,default = 1000,)
|
||
parse.add_argument('--warmup_steps',dest = 'warmup_steps',type = int,default = 1000,)
|
||
parse.add_argument('--mode',dest = 'mode',type = str,default = 'train',)
|
||
parse.add_argument('--ckpt',dest = 'ckpt',type = str,default = None,)
|
||
parse.add_argument('--respath',dest = 'respath',type = str,default='./model_save')
|
||
parse.add_argument('--backbone',dest = 'backbone',type = str,default = 'STDCNet813')
|
||
parse.add_argument('--pretrain_path',dest = 'pretrain_path',type = str,default='./checkpoints2/STDCNet813M_73.91.tar',)
|
||
parse.add_argument('--use_conv_last',dest = 'use_conv_last',type = str2bool,default = False,)
|
||
parse.add_argument('--use_boundary_2',dest = 'use_boundary_2',type = str2bool,default = False,)
|
||
parse.add_argument('--use_boundary_4',dest = 'use_boundary_4',type = str2bool,default = False,)
|
||
parse.add_argument('--use_boundary_8',dest = 'use_boundary_8',type = str2bool,default = True)
|
||
parse.add_argument('--use_boundary_16',dest = 'use_boundary_16',type = str2bool,default = False,)
|
||
return parse.parse_args()
|
||
|
||
|
||
def train():
|
||
|
||
|
||
|
||
args = parse_args()
|
||
with open(args.parJson.strip(),'r') as fp:
|
||
par=json.load(fp)
|
||
#print('args.gpuId.strip :',args.gpuId.strip())
|
||
#os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpuId.strip())
|
||
|
||
|
||
torch.cuda.set_device(int(args.gpuId.strip()))
|
||
#print('line72:',torch.cuda.current_device())
|
||
save_pth_path = os.path.join(args.respath, 'pths')
|
||
#dspth = '../../data/CarRoadLane/'
|
||
dspth = par['dspth']
|
||
|
||
if not osp.exists(save_pth_path):
|
||
os.makedirs(save_pth_path)
|
||
|
||
#torch.cuda.set_device(args.local_rank)
|
||
########################################################################fenbushi
|
||
|
||
setup_logger(args.respath)
|
||
|
||
#n_classes = 4 # 改动
|
||
n_classes = par['n_classes']
|
||
|
||
n_img_per_gpu = args.n_img_per_gpu
|
||
n_workers_train = args.n_workers_train
|
||
n_workers_val = args.n_workers_val
|
||
use_boundary_16 = args.use_boundary_16
|
||
use_boundary_8 = args.use_boundary_8
|
||
use_boundary_4 = args.use_boundary_4
|
||
use_boundary_2 = args.use_boundary_2
|
||
|
||
mode = args.mode # train
|
||
#cropsize = [640, 360]
|
||
cropsize = eval(par['cropsize'])
|
||
#labelJson='./heliushuju_info.json'
|
||
labelJson = par['labelJson']
|
||
randomscale = (0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0, 1.125, 1.25, 1.375, 1.5)
|
||
#print('line102:',torch.cuda.current_device())
|
||
|
||
|
||
ds = Heliushuju(dspth, cropsize=cropsize, mode=mode,labelJson=labelJson, randomscale=randomscale)
|
||
sampler = None
|
||
# #################################################################################################fenbushi
|
||
# sampler = torch.utils.data.distributed.DistributedSampler(ds)
|
||
dl = DataLoader(ds,
|
||
batch_size = n_img_per_gpu,
|
||
shuffle = False,
|
||
sampler = sampler,
|
||
num_workers = n_workers_train,
|
||
pin_memory = False,
|
||
drop_last = True)
|
||
dsval = Heliushuju(dspth, mode='val',labelJson=labelJson, randomscale=randomscale)
|
||
sampler_val = None
|
||
dlval = DataLoader(dsval,
|
||
batch_size = 1,
|
||
shuffle = False,
|
||
sampler = sampler_val,
|
||
num_workers = n_workers_val,
|
||
drop_last = False)
|
||
#print('line124:',torch.cuda.current_device())
|
||
## model
|
||
#ignore_idx = 255
|
||
ignore_idx = par['ignore_idx']
|
||
net = BiSeNet(backbone=args.backbone, n_classes=n_classes, pretrain_model=args.pretrain_path,
|
||
use_boundary_2=use_boundary_2, use_boundary_4=use_boundary_4, use_boundary_8=use_boundary_8,
|
||
use_boundary_16=use_boundary_16, use_conv_last=args.use_conv_last)
|
||
|
||
if not args.ckpt is None:
|
||
net.load_state_dict(torch.load(args.ckpt, map_location='cpu'))
|
||
print('load checkpoints model:',args.ckpt)
|
||
|
||
#print('line136:',torch.cuda.current_device())
|
||
net.cuda()
|
||
net.train()
|
||
#print('line142:',net.device)
|
||
#net = nn.DataParallel(net, device_ids=[0])###########################################################################
|
||
#print('line144:',net.device)
|
||
score_thres = 0.7
|
||
n_min = n_img_per_gpu*cropsize[0]*cropsize[1]//16
|
||
criteria_p = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
|
||
criteria_16 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
|
||
criteria_32 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
|
||
boundary_loss_func = DetailAggregateLoss()
|
||
## optimizer
|
||
maxmIOU50 = 0.
|
||
maxmIOU75 = 0.
|
||
momentum = 0.9
|
||
weight_decay = 5e-4
|
||
lr_start = 1e-2
|
||
max_iter = args.max_iter
|
||
save_iter_sep = args.save_iter_sep
|
||
power = 0.9
|
||
warmup_steps = args.warmup_steps
|
||
warmup_start_lr = 1e-5
|
||
print('max_iter: ', max_iter)
|
||
print('save_iter_sep: ', save_iter_sep)
|
||
print('warmup_steps: ', warmup_steps)
|
||
optim = Optimizer(
|
||
#model = net.module,
|
||
model = net,
|
||
loss = boundary_loss_func,
|
||
lr0 = lr_start,
|
||
momentum = momentum,
|
||
wd = weight_decay,
|
||
warmup_steps = warmup_steps,
|
||
warmup_start_lr = warmup_start_lr,
|
||
max_iter = max_iter,
|
||
power = power)
|
||
|
||
## train loop
|
||
msg_iter = 50
|
||
loss_avg = []
|
||
loss_boundery_bce = []
|
||
loss_boundery_dice = []
|
||
st = glob_st = time.time()
|
||
diter = iter(dl)
|
||
# diter = enumerate(dl)
|
||
epoch = 0
|
||
for it in range(max_iter):
|
||
try:
|
||
im, lb = diter.__next__()
|
||
# print(im.size()[0])
|
||
# im, lb = next(diter)
|
||
if not im.size()[0]==n_img_per_gpu: raise StopIteration
|
||
except StopIteration:
|
||
epoch += 1
|
||
# sampler.set_epoch(epoch)
|
||
diter = iter(dl)
|
||
im, lb = next(diter)
|
||
im = im.cuda()
|
||
lb = lb.cuda()
|
||
H, W = im.size()[2:]
|
||
|
||
lb = torch.squeeze(lb, 1) # lb.shape : torch.Size([8, 360, 640])
|
||
|
||
# lb = torch.argmax(lb, dim=3) # 添加(训练高速路时，需要添加这行代码，训练水域分割时，将这行代码注释掉)
|
||
|
||
optim.zero_grad()
|
||
|
||
|
||
if use_boundary_2 and use_boundary_4 and use_boundary_8:
|
||
out, out16, out32, detail2, detail4, detail8 = net(im)
|
||
|
||
if (not use_boundary_2) and use_boundary_4 and use_boundary_8:
|
||
out, out16, out32, detail4, detail8 = net(im)
|
||
|
||
if (not use_boundary_2) and (not use_boundary_4) and use_boundary_8:#######True
|
||
out, out16, out32, detail8 = net(im)
|
||
|
||
if (not use_boundary_2) and (not use_boundary_4) and (not use_boundary_8):
|
||
out, out16, out32 = net(im)
|
||
|
||
lossp = criteria_p(out, lb)
|
||
loss2 = criteria_16(out16, lb)
|
||
loss3 = criteria_32(out32, lb)
|
||
#print(' line220 : img.size:', H,W,'label.size：',lb.size(),' pred size:',out.size() )
|
||
boundery_bce_loss = 0.
|
||
boundery_dice_loss = 0.
|
||
|
||
if use_boundary_2:
|
||
boundery_bce_loss2, boundery_dice_loss2 = boundary_loss_func(detail2, lb)
|
||
boundery_bce_loss += boundery_bce_loss2
|
||
boundery_dice_loss += boundery_dice_loss2
|
||
|
||
if use_boundary_4:
|
||
boundery_bce_loss4, boundery_dice_loss4 = boundary_loss_func(detail4, lb)
|
||
boundery_bce_loss += boundery_bce_loss4
|
||
boundery_dice_loss += boundery_dice_loss4
|
||
|
||
if use_boundary_8:######
|
||
boundery_bce_loss8, boundery_dice_loss8 = boundary_loss_func(detail8, lb)
|
||
boundery_bce_loss += boundery_bce_loss8
|
||
boundery_dice_loss += boundery_dice_loss8
|
||
|
||
loss = lossp + loss2 + loss3 + boundery_bce_loss + boundery_dice_loss
|
||
|
||
loss.backward()
|
||
optim.step()
|
||
|
||
loss_avg.append(loss.item())
|
||
|
||
loss_boundery_bce.append(boundery_bce_loss.item())
|
||
loss_boundery_dice.append(boundery_dice_loss.item())
|
||
|
||
## print training log message
|
||
if (it+1)%msg_iter==0:
|
||
loss_avg = sum(loss_avg) / len(loss_avg)
|
||
lr = optim.lr
|
||
ed = time.time()
|
||
t_intv, glob_t_intv = ed - st, ed - glob_st
|
||
eta = int((max_iter - it) * (glob_t_intv / it))
|
||
eta = str(datetime.timedelta(seconds=eta))
|
||
|
||
loss_boundery_bce_avg = sum(loss_boundery_bce) / len(loss_boundery_bce)
|
||
loss_boundery_dice_avg = sum(loss_boundery_dice) / len(loss_boundery_dice)
|
||
msg = ', '.join([
|
||
'it: {it}/{max_it}',
|
||
'lr: {lr:4f}',
|
||
'loss: {loss:.4f}',
|
||
'boundery_bce_loss: {boundery_bce_loss:.4f}',
|
||
'boundery_dice_loss: {boundery_dice_loss:.4f}',
|
||
'eta: {eta}',
|
||
'time: {time:.4f}',
|
||
]).format(
|
||
it = it+1,
|
||
max_it = max_iter,
|
||
lr = lr,
|
||
loss = loss_avg,
|
||
boundery_bce_loss = loss_boundery_bce_avg,
|
||
boundery_dice_loss = loss_boundery_dice_avg,
|
||
time = t_intv,
|
||
eta = eta
|
||
)
|
||
|
||
logger.info(msg)
|
||
loss_avg = []
|
||
loss_boundery_bce = []
|
||
loss_boundery_dice = []
|
||
st = ed
|
||
# print(boundary_loss_func.get_params())
|
||
if (it+1)%save_iter_sep==0:# and it != 0:
|
||
#if (it+1)%50==0:# and it != 0:
|
||
|
||
## model
|
||
logger.info('evaluating the model ...')
|
||
logger.info('setup and restore model')
|
||
|
||
net.eval()
|
||
|
||
# ## evaluator
|
||
logger.info('compute the mIOU')
|
||
with torch.no_grad():
|
||
single_scale1 = MscEvalV0()
|
||
mIOU50 = single_scale1(net, dlval, n_classes)
|
||
|
||
single_scale2 = MscEvalV0(scale=0.75)
|
||
mIOU75 = single_scale2(net, dlval, n_classes)
|
||
|
||
save_pth = osp.join(save_pth_path, 'model_iter{}_mIOU50_{}_mIOU75_{}.pth'
|
||
.format(it+1, str(round(mIOU50, 4)), str(round(mIOU75, 4))))
|
||
|
||
state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict()
|
||
# if dist.get_rank()==0:
|
||
torch.save(state, save_pth)
|
||
|
||
logger.info('training iteration {}, model saved to: {}'.format(it+1, save_pth))
|
||
|
||
if mIOU50 > maxmIOU50:
|
||
maxmIOU50 = mIOU50
|
||
save_pth = osp.join(save_pth_path, 'model_maxmIOU50.pth'.format(it+1))
|
||
state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict()
|
||
# if dist.get_rank()==0:
|
||
torch.save(state, save_pth)
|
||
|
||
logger.info('max mIOU model saved to: {}'.format(save_pth))
|
||
|
||
if mIOU75 > maxmIOU75:
|
||
maxmIOU75 = mIOU75
|
||
save_pth = osp.join(save_pth_path, 'model_maxmIOU75.pth'.format(it+1))
|
||
state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict()
|
||
# if dist.get_rank()==0: torch.save(state, save_pth)
|
||
torch.save(state, save_pth)
|
||
logger.info('max mIOU model saved to: {}'.format(save_pth))
|
||
|
||
logger.info('mIOU50 is: {}, mIOU75 is: {}'.format(mIOU50, mIOU75))
|
||
logger.info('maxmIOU50 is: {}, maxmIOU75 is: {}.'.format(maxmIOU50, maxmIOU75))
|
||
|
||
net.train()
|
||
|
||
## dump the final model
|
||
save_pth = osp.join(save_pth_path, 'model_final.pth')
|
||
net.cpu()
|
||
state = net.module.state_dict() if hasattr(net, 'module') else net.state_dict()
|
||
# if dist.get_rank()==0: torch.save(state, save_pth)
|
||
torch.save(state, save_pth)
|
||
logger.info('training done, model saved to: {}'.format(save_pth))
|
||
print('epoch: ', epoch)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
train()
|