def train():
    args = parse_args()
    save_pth_path = os.path.join(args.respath, 'pths')
    dspth = './data'
    if (not osp.exists(save_pth_path)):
        os.makedirs(save_pth_path)
    torch.cuda.set_device(args.local_rank)
    dist.init_process_group(backend='nccl', init_method='env://', world_size=torch.cuda.device_count(), rank=args.local_rank)
    setup_logger(args.respath)
    n_classes = 19
    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
    cropsize = [1024, 512]
    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)
    if (dist.get_rank() == 0):
        logger.info('n_workers_train: {}'.format(n_workers_train))
        logger.info('n_workers_val: {}'.format(n_workers_val))
        logger.info('use_boundary_2: {}'.format(use_boundary_2))
        logger.info('use_boundary_4: {}'.format(use_boundary_4))
        logger.info('use_boundary_8: {}'.format(use_boundary_8))
        logger.info('use_boundary_16: {}'.format(use_boundary_16))
        logger.info('mode: {}'.format(args.mode))
    ds = CityScapes(dspth, cropsize=cropsize, mode=mode, randomscale=randomscale)
    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 = CityScapes(dspth, mode='val', randomscale=randomscale)
    sampler_val = torch.utils.data.distributed.DistributedSampler(dsval)
    dlval = DataLoader(dsval, batch_size=2, shuffle=False, sampler=sampler_val, num_workers=n_workers_val, drop_last=False)
    ignore_idx = 255
    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'))
    net.cuda()
    net.train()
    net = nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)
    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()
    maxmIOU50 = 0.0
    maxmIOU75 = 0.0
    momentum = 0.9
    weight_decay = 0.0005
    lr_start = 0.01
    max_iter = args.max_iter
    save_iter_sep = args.save_iter_sep
    power = 0.9
    warmup_steps = args.warmup_steps
    warmup_start_lr = 1e-05
    if (dist.get_rank() == 0):
        print('max_iter: ', max_iter)
        print('save_iter_sep: ', save_iter_sep)
        print('warmup_steps: ', warmup_steps)
    optim = Optimizer(model=net.module, 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)
    msg_iter = 50
    loss_avg = []
    loss_boundery_bce = []
    loss_boundery_dice = []
    st = glob_st = time.time()
    diter = iter(dl)
    epoch = 0
    for it in range(max_iter):
        try:
            (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)
        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):
            (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)
        boundery_bce_loss = 0.0
        boundery_dice_loss = 0.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())
        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
        if (((it + 1) % save_iter_sep) == 0):
            logger.info('evaluating the model ...')
            logger.info('setup and restore model')
            net.eval()
            logger.info('compute the mIOU')
            with torch.no_grad():
                single_scale1 = MscEvalV0()
                mIOU50 = single_scale1(net, dlval, n_classes)
            save_pth = osp.join(save_pth_path, 'model_iter{}_mIOU50_{}.pth'.format((it + 1), str(round(mIOU50, 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))
            logger.info('mIOU50 is: {}'.format(mIOU50))
            logger.info('maxmIOU50 is: {}'.format(maxmIOU50))
            net.train()
    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)
    logger.info('training done, model saved to: {}'.format(save_pth))
    print('epoch: ', epoch)
def train():
    args = parse_args()
    save_pth_path = os.path.join(args.respath, 'pths')
    dspth = './data'
    if (not osp.exists(save_pth_path)):
        os.makedirs(save_pth_path)
    torch.cuda.set_device(args.local_rank)
    dist.init_process_group(backend='nccl', init_method='env://', world_size=torch.cuda.device_count(), rank=args.local_rank)
    setup_logger(args.respath)
    n_classes = 19
    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
    cropsize = [1024, 512]
    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)
    if (dist.get_rank() == 0):
        logger.info('n_workers_train: {}'.format(n_workers_train))
        logger.info('n_workers_val: {}'.format(n_workers_val))
        logger.info('use_boundary_2: {}'.format(use_boundary_2))
        logger.info('use_boundary_4: {}'.format(use_boundary_4))
        logger.info('use_boundary_8: {}'.format(use_boundary_8))
        logger.info('use_boundary_16: {}'.format(use_boundary_16))
        logger.info('mode: {}'.format(args.mode))
    ds = CityScapes(dspth, cropsize=cropsize, mode=mode, randomscale=randomscale)
    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 = CityScapes(dspth, mode='val', randomscale=randomscale)
    sampler_val = torch.utils.data.distributed.DistributedSampler(dsval)
    dlval = DataLoader(dsval, batch_size=2, shuffle=False, sampler=sampler_val, num_workers=n_workers_val, drop_last=False)
    ignore_idx = 255
    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'))
    net.cuda()
    net.train()
    net = nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)
    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()
    maxmIOU50 = 0.0
    maxmIOU75 = 0.0
    momentum = 0.9
    weight_decay = 0.0005
    lr_start = 0.01
    max_iter = args.max_iter
    save_iter_sep = args.save_iter_sep
    power = 0.9
    warmup_steps = args.warmup_steps
    warmup_start_lr = 1e-05
    if (dist.get_rank() == 0):
        print('max_iter: ', max_iter)
        print('save_iter_sep: ', save_iter_sep)
        print('warmup_steps: ', warmup_steps)
    optim = Optimizer(model=net.module, 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)
    msg_iter = 50
    loss_avg = []
    loss_boundery_bce = []
    loss_boundery_dice = []
    st = glob_st = time.time()
    diter = iter(dl)
    epoch = 0
    for it in range(max_iter):
        try:
            (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)
        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):
            (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)
        boundery_bce_loss = 0.0
        boundery_dice_loss = 0.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())
        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
        if (((it + 1) % save_iter_sep) == 0):
            logger.info('evaluating the model ...')
            logger.info('setup and restore model')
            net.eval()
            logger.info('compute the mIOU')
            with torch.no_grad():
                single_scale1 = MscEvalV0()
                mIOU50 = single_scale1(net, dlval, n_classes)
            save_pth = osp.join(save_pth_path, 'model_iter{}_mIOU50_{}.pth'.format((it + 1), str(round(mIOU50, 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))
            logger.info('mIOU50 is: {}'.format(mIOU50))
            logger.info('maxmIOU50 is: {}'.format(maxmIOU50))
            net.train()
    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)
    logger.info('training done, model saved to: {}'.format(save_pth))
    print('epoch: ', epoch)