def get_tbsm()

in tbsm_pytorch.py [0:0]

• 78 lines of code
• 18 McCabe index (conditional complexity)

def get_tbsm(args, use_gpu):

    # train, test, or train-test
    modes = args.mode.split("-")
    model_file = args.save_model

    if args.debug_mode:
        print("model_file: ", model_file)
        print("model_type: ", args.model_type)

    if use_gpu:
        ngpus = torch.cuda.device_count()  # 1
        devicenum = "cuda:" + str(args.device_num % ngpus)
        print("device:", devicenum)
        device = torch.device(devicenum)
        print("Using {} GPU(s)...".format(ngpus))
    else:
        device = torch.device("cpu")
        print("Using CPU...")

    # prepare dlrm arch
    m_spa = args.arch_sparse_feature_size
    # this is an array of sizes of cat features
    ln_emb = np.fromstring(args.arch_embedding_size, dtype=int, sep="-")
    num_fea = ln_emb.size + 1  # user: num sparse + bot_mlp(all dense)
    ln_bot = np.fromstring(args.arch_mlp_bot, dtype=int, sep="-")
    #  m_den = ln_bot[0]
    ln_bot[ln_bot.size - 1] = m_spa  # enforcing
    m_den_out = ln_bot[ln_bot.size - 1]  # must be == m_spa (embed dim)

    if args.arch_interaction_op == "dot":
        # approach 1: all
        # num_int = num_fea * num_fea + m_den_out
        # approach 2: unique
        if args.arch_interaction_itself:
            num_int = (num_fea * (num_fea + 1)) // 2 + m_den_out
        else:
            num_int = (num_fea * (num_fea - 1)) // 2 + m_den_out
    elif args.arch_interaction_op == "cat":
        num_int = num_fea * m_den_out
    else:
        sys.exit(
            "ERROR: --arch-interaction-op="
            + args.arch_interaction_op
            + " is not supported"
        )
    arch_mlp_top_adjusted = str(num_int) + "-" + args.arch_mlp_top
    ln_top = np.fromstring(arch_mlp_top_adjusted, dtype=int, sep="-")
    # sigmoid_top = len(ln_top) - 2    # used only if length_ts == 1
    # attention mlp (will be automatically adjusted so that first and last
    # layer correspond to number of vectors (ts_length) used in attention)
    ln_atn = np.fromstring(args.tsl_mlp, dtype=int, sep="-")
    # context MLP (with automatically adjusted first layer)
    if args.model_type == "mha":
        num_cat = (int(args.mha_num_heads) + 1) * ln_top[-1]    # mha with k heads + w
    else:         # tsl or rnn
        num_cat = 2 * ln_top[-1]   # [c,w]
    arch_mlp_adjusted = str(num_cat) + "-" + args.arch_mlp
    ln_mlp = np.fromstring(arch_mlp_adjusted, dtype=int, sep="-")

    # construct TBSM
    tbsm = TBSM_Net(
        m_spa,
        ln_emb,
        ln_bot,
        ln_top,
        args.arch_interaction_op,
        args.arch_interaction_itself,
        ln_mlp,
        ln_atn,
        args.tsl_interaction_op,
        args.tsl_mechanism,
        args.ts_length,
        -1,
        args.model_type,
        args.tsl_seq,
        args.tsl_proj,
        args.tsl_inner,
        args.tsl_num_heads,
        args.mha_num_heads,
        args.rnn_num_layers,
        args.debug_mode,
    )

    # move model to gpu
    if use_gpu:
        tbsm = tbsm.to(device)  # .cuda()

    # load existing pre-trained model if needed
    if path.exists(model_file):
        if modes[0] == "test" or (len(modes) > 1 and modes[1] == "test"):
            if use_gpu:
                ld_model = torch.load(
                    model_file,
                    map_location=torch.device('cuda')
                )
            else:
                # when targeting inference on CPU
                ld_model = torch.load(model_file, map_location=torch.device('cpu'))

            tbsm.load_state_dict(ld_model['model_state_dict'])

    return tbsm, device