config.py

from install import cloud_prefix
from pathlib import Path
import lz
from lz import *
from torch.nn import CrossEntropyLoss
from tools.vat import VATLoss
from torchvision import transforms as trans

# torch.autograd.set_detect_anomaly(True)
# print = lambda x: logging.info(f'do not prt {x}')
debug = os.environ.get('debug', '')
if cloud_prefix.startswith('s3:'):
    debug = False
dist = False
num_devs = torch.cuda.device_count() if not debug else 1
assert num_devs != 0, f'{num_devs} {debug}'
if dist:
    num_devs = 1

conf = edict()
try:
    import mox_patch_0603_v4
    import moxing.pytorch as mox

    conf.cloud = True
except:
    logging.info('not in the cloud')
    conf.cloud = False
conf.num_workers = 4 * num_devs if not debug else 0
conf.num_devs = num_devs
conf.no_eval = False
conf.start_eval = False
conf.loss = 'arcface'  # adamarcface adamrg adacos softmax arcface arcfaceneg cosface arcsinmrg

conf.writer = None
conf.local_rank = None
conf.num_clss = None
conf.dop = None  # top_imp
conf.id2range_dop = None  # sub_imp
conf.explored = None

conf.hostname = hostname()

if "amax" in conf.hostname:
    conf.data_path = Path('/data2/share/')
elif "g5500" in conf.hostname:
    conf.data_path = Path('/home/zl/zl_data/')
else:
    assert conf.cloud
    conf.data_path = Path(f'{cloud_prefix}/zl/datasets/')
    print('hostname is ', hostname)
    conf.cloud = True
conf.work_path = Path(f'work_space/r100.bs.bak')
conf.load_from = cloud_prefix + '/work_space/alphaf64.dop.tri.chsfst.ft.4/models'  # todo
# conf.load_from = './work_space/r100.bs.2/models'
if conf.cloud:
    conf.work_path = Path(f'{cloud_prefix}') / conf.work_path

conf.model_path = conf.work_path / 'models'
conf.log_path = conf.work_path / 'log'
conf.save_path = conf.work_path / 'save'
vgg_folder = conf.data_path / 'faces_vgg_112x112'
ms1m_folder = conf.data_path / 'faces_ms1m_112x112'
glint_folder = conf.data_path / 'glint'
emore_folder = conf.data_path / 'faces_emore'
asia_emore = conf.data_path / 'asia_emore'
glint_test = conf.data_path / 'glint_test'
alpha_f64 = conf.data_path / 'alpha_f64'
alpha_jk = conf.data_path / 'alpha_jk'
# casia_folder = conf.data_path / 'casia'  # the cleaned one todo may need the other for exploring the noise
retina_folder = conf.data_path / 'ms1m-retinaface-t1'
dingyi_folder = conf.data_path / 'faces_casia'

conf.use_data_folder = alpha_f64  # if not debug else alpha_jk # for test
conf.dataset_name = str(conf.use_data_folder).split('/')[-1]
conf.clean_ids = None  # np.asarray(msgpack_load(root_path + 'train.configs/noise.20.pk', allow_np=False))

if conf.use_data_folder == ms1m_folder:
    conf.cutoff = 0
elif conf.use_data_folder == glint_folder:
    conf.cutoff = 15
elif conf.use_data_folder == emore_folder:
    conf.cutoff = 0
elif conf.use_data_folder == asia_emore:
    conf.cutoff = 10
else:
    conf.cutoff = 0
conf.mining = 'rand.id'  # todo balance opt # dop rand.id # todo imp.grad imp.loss imp rand.img(slow)
conf.mining_init = 1 if conf.mining == 'rand.id' else -1  # imp 1.6; rand.id 1; dop -1
conf.rand_ratio = 6 / 27

conf.margin = .5  # todo do not forget if use adacos!
conf.margin2 = .2
conf.topk = 3
conf.fgg = ''  # g gg ''
conf.fgg_wei = 0.  # 1
conf.tri_wei = 0.
conf.tri_scale = 64
conf.scale = 64  # 48 64
conf.instances = 4
conf.pos_wei = 1.  # 2
conf.neg_wei = 1.  # 50

conf.phi = 1.9
conf.input_rg_255 = False
conf.input_size = 112  # 128 224 112
conf.embedding_size = 512  # 2048#
conf.drop_ratio = .4
conf.conv2dmask_drop_ratio = .2
conf.lambda_runtime_reg = 5
conf.net_mode = 'ir_se'  # effnet mbfc sglpth hrnet mbv3 mobilefacenet ir_se resnext densenet widerresnet
conf.decs = None
conf.net_depth = 100  # 100 121 169 201 264 50 20
conf.mb_mode = 'face.large'
conf.mb_mult = 1.285
# conf.mb_mode = 'face.small'
# conf.mb_mult = 2.005 # 1.37
conf.mbfc_wm = 1  # 1.2 ** conf.phi
conf.mbfc_dm = 2  # 1.56 ** conf.phi
conf.mbfc_se = True
conf.lpf = False
conf.eff_name = 'efficientnet-b0'

conf.test_transform = trans.Compose([
    trans.ToTensor(),
    trans.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
conf.use_loader = 'torch'  # todo mxent lack val path and testing speed
conf.flip = True

conf.upgrade_irse = True
conf.upgrade_bnneck = False  # todo may pretrain by imgnet
conf.use_redis = False
conf.use_chkpnt = False
conf.chs_first = True
conf.prof = False
conf.fast_load = False
conf.ipabn = True
conf.cvt_ipabn = False  # todo
conf.kd = False
conf.sftlbl_from_file = True
conf.alpha = .95
conf.temperature = 24
conf.teacher_head_dev = 0  # num_devs - 1  # -1 #
conf.teacher_head_in_dloader = False  # todo bug when True

conf.online_imp = False
conf.use_test = False  # 'ijbc' 'glint' False 'cfp_fp'
conf.model1_dev = list(range(num_devs))
conf.model2_dev = list(range(num_devs))
conf.tau = 0.05
conf.mutual_learning = 0

conf.fp16 = True
conf.opt_level = "O1"
if debug:
    conf.batch_size = 32 * num_devs
else:
    if conf.cloud:
        conf.batch_size = 224 * num_devs
    else:
        conf.batch_size = 64 * num_devs

conf.ftbs_mult = 2
conf.board_loss_every = 99
conf.log_interval = 999 if not debug else 99
conf.need_tb = True
conf.other_every = None if not debug else 99
conf.num_recs = 1
conf.acc_grad = 1
# --------------------Training Config ------------------------
conf.weight_decay = 4e-5  # 4e-5 5e-4, 1e-6 for 1e-3, 0.3 for 3e-3
conf.use_opt = 'sgd'  # adabound adam radam ranger sgd
conf.adam_betas1 = .95  # .85 to .95
conf.adam_betas2 = .999  # 0.999 0.99
conf.final_lr = 1e-1
conf.lr = 1e-3  # 1e-1  #
conf.lr_gamma = 0.1
conf.lr_mult = 1
conf.wd_mult = 10
conf.start_epoch = 0
conf.start_step = 0
# conf.epochs = 37
# conf.milestones = (23,32)
conf.schedule = 'mstep'  # cosanl cycle
conf.cycle_up = 0.
conf.cycle_down = 1.
conf.cycle_gamma = .7
conf.cycle_max_lr_mult = 10.
conf.cycle_momentum=False
conf.epochs = 5  # 18
conf.milestones = (2,)  # (9,13)
conf.warmup = 0  # conf.epochs/25 # 1 0
conf.epoch_less_iter = 1
conf.momentum = 0.9
conf.pin_memory = True
conf.fill_cache = .7 if not debug else 0
conf.val_ijbx = False
conf.spec_norm = False
conf.use_of = False
conf.use_act = "prelu"
conf.bottle_neck = False
conf.never_stop = False
conf.n_sma = 4
conf.out_type = 'fc'
conf.mid_type = ''  # 'gpool'  # 'fc'
conf.use_bl = False
conf.arch_ft = False  # maybe this improves
conf.pfe = False
conf.ds = False
conf.use_in = False
conf.sigmoid_mult = 1  # may better
conf.prev_wei = 0.9
conf.head_parallel = ''
conf.head_init = ''#all
conf.head_init_zjjk=False
conf.head_load = False
conf.swa = '' # modify


# todo may use kl_div to speed up
class CrossEntropyLabelSmooth(nn.Module):
    """Cross entropy loss with label smoothing regularizer.
    Reference:
    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
    Equation: y = (1 - epsilon) * y + epsilon / K.
    Args:
        num_classes (int): number of classes.
        epsilon (float): weight.
    """

    def __init__(self, epsilon=0.1, ):
        super(CrossEntropyLabelSmooth, self).__init__()
        self.epsilon = epsilon
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self, inputs, targets):
        """
        Args:
            inputs: prediction matrix (before softmax) with shape (bs, num_classes)
            targets: ground truth labels with shape (num_classes)
        """
        log_probs = self.logsoftmax(inputs)
        targets1 = torch.zeros(log_probs.size()).scatter_(
            1, targets.unsqueeze(1).data.cpu(), 1).cuda()
        # targets2 = torch.cuda.FloatTensor(inputs.size()).fill_(0).scatter_(1, targets.unsqueeze(1).detach(), 1)
        targets3 = (1 - self.epsilon) * targets1 + \
                   self.epsilon / inputs.shape[1]
        loss = (-targets3 * log_probs).mean(0).sum()
        return loss


class CrossEntropySigSoft(nn.Module):

    def forward(self, inputs, targets):
        log_probs = logsigsoftmax(inputs)
        bs = inputs.shape[0]
        idx_ = torch.arange(0, bs, dtype=torch.long)
        log_probs2 = log_probs[idx_, targets]
        loss = - log_probs2.mean()
        return loss


def logsigsoftmax(logits):
    """
    Computes sigsoftmax from the paper - https://arxiv.org/pdf/1805.10829.pdf
    """
    max_values = torch.max(logits, 1, keepdim=True)[0]
    exp_logits_sigmoided = torch.exp(logits - max_values) * torch.sigmoid(logits)
    sum_exp_logits_sigmoided = exp_logits_sigmoided.sum(1, keepdim=True)
    log_probs = logits - max_values + F.logsigmoid(logits) - torch.log(sum_exp_logits_sigmoided)
    return log_probs


class CrossEntropySigSoft2(nn.Module):

    def forward(self, inputs, targets):
        inputs2 = inputs + F.logsigmoid(inputs / 10)  # todo
        loss = F.cross_entropy(inputs2, targets)
        return loss


conf.ce_loss = CrossEntropyLoss()  # CrossEntropySigSoft2()  #   CrossEntropyLabelSmooth()
if conf.use_test:
    conf.vat_loss_func = VATLoss(xi=1e-6, eps=8, ip=1)
conf.need_log = True
conf.batch_size = conf.batch_size // conf.instances * conf.instances