[WWW2023] PyTorch implementation of DUET: A Tuning-Free Device-Cloud Collaborative Parameters Generation Framework for Efficient Device Model Generalization.
# coding=utf-8
import os
import sys
import logging

import torch
# import common_io
import pandas as pd


logger = logging.getLogger(__name__)


class LoaderBase(torch.utils.data.IterableDataset):
    def __init__(self, table_name, slice_id, slice_count, columns, is_train):
        super(LoaderBase).__init__()
        self._table_name = table_name
        self._slice_id_init = slice_id
        self._slice_count = slice_count
        self._is_train = is_train
        self._columns = columns
        print("dataset init")
        self._line_count = self._get_dataset_info()

        if self._is_train:
            self._repeat = 999999999
        else:
            self._repeat = 0

    def parse_data(self, data):
        raise NotImplementedError

    def batchify(self):
        raise NotImplementedError

    def _get_dataset_info(self):
        _line_count = 0
        with open(self._table_name, "r") as reader:
            for line in reader:
                _line_count += 1
        print("table name: ", self._table_name)
        print("_line_count: ", _line_count)
        return _line_count

    def _reopen_reader(self, original_reader):
        if original_reader is not None:
            original_reader.close()
        reader = open(self._table_name, "r")
        print("_reopen_reader")
        self._get_dataset_info()
        return reader

    def __iter__(self):
        with open(self._table_name, "r") as reader:
            for data in reader:
                sample = self.parse_data(data)
                yield sample
                
loader/loader_base.py
HelloZicky/DUET
import numpy as np
import torch

from loader.loader_base import LoaderBase

from util import consts


class MetaSequenceDataLoader(LoaderBase):
    def __init__(self, table_name, slice_id, slice_count, is_train):
        super(MetaSequenceDataLoader, self).__init__(
            table_name=table_name,
            slice_id=slice_id,
            slice_count=slice_count,
            columns=[consts.FIELD_USER_ID, consts.FIELD_TARGET_ID, consts.FIELD_TRIGGER_SEQUENCE,
                     consts.FIELD_CLK_SEQUENCE, consts.FIELD_LABEL],
            is_train=is_train
        )

    def parse_data(self, data):
        data = data.strip("\n").split(";")
        return (
            int(data[0]),
            int(data[1]),
            np.fromstring(data[2], dtype=np.int32, sep=","),
            np.fromstring(data[3], dtype=np.int32, sep=","),
            float(data[4])
        )

    @staticmethod
    def batchify(data):
        return {
            consts.FIELD_USER_ID: torch.from_numpy(np.stack([item[0] for item in data], axis=0)),
            consts.FIELD_TARGET_ID: torch.from_numpy(np.stack([item[1] for item in data], axis=0)),
            consts.FIELD_TRIGGER_SEQUENCE: torch.from_numpy(np.stack([item[2] for item in data], axis=0)),
            consts.FIELD_CLK_SEQUENCE: torch.from_numpy(np.stack([item[3] for item in data], axis=0)),
            consts.FIELD_LABEL: torch.from_numpy(np.stack([item[4] for item in data], axis=0))
        }


loader/multi_metric_meta_sequence_dataloader.py
HelloZicky/DUET
# coding=utf-8
import os
import time
import json
import logging
import math
import argparse
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import torch
import torch.multiprocessing as mp
from torch import nn

import model
from util.timer import Timer
from util import args_processing as ap
from util import consts
from util import env
from util import new_metrics
from loader import multi_metric_meta_sequence_dataloader as meta_sequence_dataloader
import numpy as np
from thop import profile

from util import utils
utils.setup_seed(0)

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tables", type=str, help="Kernels configuration for CNN")
    parser.add_argument("--bucket", type=str, default=None, help="Bucket name for external storage")
    parser.add_argument("--dataset", type=str, default="alipay", help="Bucket name for external storage")

    parser.add_argument("--max_steps", type=int, help="Number of iterations before stopping")
    parser.add_argument("--snapshot", type=int, help="Number of iterations to dump model")
    parser.add_argument("--checkpoint_dir", type=str, help="Path of the checkpoint path")
    parser.add_argument("--learning_rate", type=str, default=0.001)
    parser.add_argument("--batch_size", type=int, default=2, help="Batch size")
    parser.add_argument("--max_epoch", type=int,  default=10, help="Max epoch")
    parser.add_argument("--num_loading_workers", type=int, default=4, help="Number of threads for loading")
    parser.add_argument("--model", type=str, help="model type")
    parser.add_argument("--init_checkpoint", type=str, default="", help="Path of the checkpoint path")
    parser.add_argument("--init_step", type=int, default=0, help="Path of the checkpoint path")

    parser.add_argument("--max_gradient_norm", type=float, default=0.)

    parser.add_argument("--arch_config_path", type=str, default=None, help="Path of model configs")
    parser.add_argument("--arch_config", type=str, default=None, help="base64-encoded model configs")

    return parser.parse_known_args()[0]


def predict(predict_dataset, model_obj, device, args, train_epoch, train_step, writer=None):

    model_obj.eval()
    model_obj.to(device)

    timer = Timer()
    log_every = 200

    pred_list = []
    y_list = []
    buffer = []
    user_id_list = []
    for step, batch_data in enumerate(predict_dataset, 1):
        logits = model_obj({
            key: value.to(device)
            for key, value in batch_data.items()
            if key not in {consts.FIELD_USER_ID, consts.FIELD_LABEL}
        })
        prob = torch.sigmoid(logits).detach().cpu().numpy()
        y = batch_data[consts.FIELD_LABEL].view(-1, 1)
        overall_auc, _, _, _ = new_metrics.calculate_overall_auc(prob, y)

        user_id_list.extend(np.array(batch_data[consts.FIELD_USER_ID].view(-1, 1)))
        pred_list.extend(prob)
        y_list.extend(np.array(y))

        buffer.extend(
            [int(user_id), float(score), float(label)]
            for user_id, score, label
            in zip(
                batch_data[consts.FIELD_USER_ID],
                prob,
                batch_data[consts.FIELD_LABEL]
            )
        )

        if step % log_every == 0:
            logger.info(
                "train_epoch={}, step={}, overall_auc={:5f}, speed={:2f} steps/s".format(
                    train_epoch, step, overall_auc, log_every / timer.tick(False)
                )
            )

    overall_auc, _, _, _ = new_metrics.calculate_overall_auc(np.array(pred_list), np.array(y_list))
    user_auc = new_metrics.calculate_user_auc(buffer)
    overall_logloss = new_metrics.calculate_overall_logloss(np.array(pred_list), np.array(y_list))
    user_ndcg5, user_hr5 = new_metrics.calculate_user_ndcg_hr(5, buffer)
    user_ndcg10, user_hr10 = new_metrics.calculate_user_ndcg_hr(10, buffer)
    user_ndcg20, user_hr20 = new_metrics.calculate_user_ndcg_hr(20, buffer)

    print("train_epoch={}, train_step={}, overall_auc={:5f}, user_auc={:5f}, overall_logloss={:5f}, "
          "user_ndcg5={:5f}, user_hr5={:5f}, user_ndcg10={:5f}, user_hr10={:5f}, user_ndcg20={:5f}, user_hr20={:5f}".
          format(train_epoch, train_step, overall_auc, user_auc, overall_logloss,
                 user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20))
    with open(os.path.join(args.checkpoint_dir, "log_ood.txt"), "a") as writer:
        print("train_epoch={}, train_step={}, overall_auc={:5f}, user_auc={:5f}, overall_logloss={:5f}, "
              "user_ndcg5={:5f}, user_hr5={:5f}, user_ndcg10={:5f}, user_hr10={:5f}, user_ndcg20={:5f}, user_hr20={:5f}".
              format(train_epoch, train_step, overall_auc, user_auc, overall_logloss,
                     user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20), file=writer)

    return overall_auc, user_auc, overall_logloss, user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20


def train(train_dataset, model_obj, device, args, pred_dataloader):
    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(
        model_obj.parameters(),
        lr=float(args.learning_rate)
    )
    model_obj.train()
    model_obj.to(device)

    print(model_obj)

    logger.info("Start training...")
    timer = Timer()
    log_every = 200
    max_step = 0
    best_auc = 0
    best_auc_ckpt_path = os.path.join(args.checkpoint_dir, "best_auc" + ".pkl")
    for epoch in range(1, args.max_epoch + 1):
        for step, batch_data in enumerate(train_dataset, 1):
            logits = model_obj({
                key: value.to(device)
                for key, value in batch_data.items()
                if key not in {consts.FIELD_USER_ID, consts.FIELD_LABEL}
            })

            loss = criterion(logits, batch_data[consts.FIELD_LABEL].view(-1, 1).to(device))
            pred, y = torch.sigmoid(logits), batch_data[consts.FIELD_LABEL].view(-1, 1)
            auc, _, _, _ = new_metrics.calculate_overall_auc(np.array(pred.detach().cpu()), np.array(y))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if step % log_every == 0:
                logger.info(
                    "epoch={}, step={}, loss={:5f}, auc={:5f}, speed={:2f} steps/s".format(
                        epoch, step, float(loss.item()), auc, log_every / timer.tick(False)
                    )
                )
            max_step = step

        pred_overall_auc, pred_user_auc, pred_overall_logloss, pred_user_ndcg5, pred_user_hr5, \
        pred_user_ndcg10, pred_user_hr10, pred_user_ndcg20, pred_user_hr20 = predict(
            predict_dataset=pred_dataloader,
            model_obj=model_obj,
            device=device,
            args=args,
            train_epoch=epoch,
            train_step=epoch * max_step,
        )
        logger.info("dump checkpoint for epoch {}".format(epoch))
        model_obj.train()
        if pred_overall_auc > best_auc:
            best_auc = pred_overall_auc
            torch.save(model_obj, best_auc_ckpt_path)


def main_worker(_):
    args = parse_args()
    ap.print_arguments(args)

    if not os.path.exists(args.checkpoint_dir):
        os.makedirs(args.checkpoint_dir)
    model_meta = model.get_model_meta(args.model)  # type: model.ModelMeta

    model_conf, raw_model_conf = ap.parse_arch_config_from_args(model_meta, args)  # type: dict

    model_obj = model_meta.model_builder(model_conf=model_conf)  # type: torch.nn.module
    print("=" * 100)
    for name, parms in model_obj.named_parameters():
        print(name)
    print("=" * 100)
    device = env.get_device()

    worker_id = worker_count = 8
    train_file, test_file = args.dataset.split(',')

    args.num_loading_workers = 1
    train_dataloader = meta_sequence_dataloader.MetaSequenceDataLoader(
        table_name=train_file,
        slice_id=0,
        slice_count=args.num_loading_workers,
        is_train=True
    )
    train_dataloader = torch.utils.data.DataLoader(
        train_dataloader,
        batch_size=args.batch_size,
        num_workers=args.num_loading_workers,
        pin_memory=True,
        collate_fn=train_dataloader.batchify,
        drop_last=False
    )

    # Setup up data loader
    pred_dataloader = meta_sequence_dataloader.MetaSequenceDataLoader(
        table_name=test_file,
        slice_id=args.num_loading_workers * worker_id,
        slice_count=args.num_loading_workers * worker_count,
        is_train=False
    )
    pred_dataloader = torch.utils.data.DataLoader(
        pred_dataloader,
        batch_size=args.batch_size,
        num_workers=args.num_loading_workers,
        pin_memory=True,
        collate_fn=pred_dataloader.batchify,
        drop_last=False
    )

    # Setup training
    train(
        train_dataset=train_dataloader,
        model_obj=model_obj,
        device=device,
        args=args,
        pred_dataloader=pred_dataloader
    )


if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    mp.spawn(main_worker, nprocs=1)


main/multi_metric_meta_train.py
HelloZicky/DUET
# coding=utf-8
import os
import time
import json
import logging
import math
import argparse
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import torch
import torch.multiprocessing as mp
from torch import nn

import model
from util.timer import Timer
from util import args_processing as ap
from util import consts
from util import env
from loader import multi_metric_meta_sequence_dataloader as sequence_dataloader
from util import new_metrics
import numpy as np
from thop import profile

from util import utils
utils.setup_seed(0)

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tables", type=str, help="Kernels configuration for CNN")
    parser.add_argument("--bucket", type=str, default=None, help="Bucket name for external storage")
    parser.add_argument("--dataset", type=str, default="alipay", help="Bucket name for external storage")

    parser.add_argument("--max_steps", type=int, help="Number of iterations before stopping")
    parser.add_argument("--snapshot", type=int, help="Number of iterations to dump model")
    parser.add_argument("--checkpoint_dir", type=str, help="Path of the checkpoint path")
    parser.add_argument("--learning_rate", type=str, default=0.001)
    parser.add_argument("--batch_size", type=int, default=2, help="Batch size")
    # parser.add_argument("--max_epoch", type=int, default=10, help="Max epoch")
    parser.add_argument("--max_epoch", type=int,  default=10, help="Max epoch")
    parser.add_argument("--num_loading_workers", type=int, default=4, help="Number of threads for loading")
    parser.add_argument("--model", type=str, help="model type")
    parser.add_argument("--init_checkpoint", type=str, default="", help="Path of the checkpoint path")
    parser.add_argument("--init_step", type=int, default=0, help="Path of the checkpoint path")

    parser.add_argument("--max_gradient_norm", type=float, default=0.)

    # If both of the two options are set, `model_config` is preferred
    parser.add_argument("--arch_config_path", type=str, default=None, help="Path of model configs")
    parser.add_argument("--arch_config", type=str, default=None, help="base64-encoded model configs")

    return parser.parse_known_args()[0]


# def predict(predict_dataset, model_obj, device, args, bucket, train_step, writer=None):
def predict(predict_dataset, model_obj, device, args, train_epoch, train_step, writer=None):
    model_obj.eval()
    model_obj.to(device)

    timer = Timer()
    log_every = 200
    pred_list = []
    y_list = []
    buffer = []
    user_id_list = []
    for step, batch_data in enumerate(predict_dataset, 1):
        logits = model_obj({
            key: value.to(device)
            for key, value in batch_data.items()
            if key not in {consts.FIELD_USER_ID, consts.FIELD_LABEL, consts.FIELD_TRIGGER_SEQUENCE}
        })

        prob = torch.sigmoid(logits).detach().cpu().numpy()
        y = batch_data[consts.FIELD_LABEL].view(-1, 1)
        overall_auc, _, _, _ = new_metrics.calculate_overall_auc(prob, y)
        user_id_list.extend(np.array(batch_data[consts.FIELD_USER_ID].view(-1, 1)))
        pred_list.extend(prob)
        y_list.extend(np.array(y))

        buffer.extend(
            [int(user_id), float(score), float(label)]
            for user_id, score, label
            in zip(
                batch_data[consts.FIELD_USER_ID],
                prob,
                batch_data[consts.FIELD_LABEL]
            )
        )
        if step % log_every == 0:
            logger.info(
                "train_epoch={}, step={}, overall_auc={:5f}, speed={:2f} steps/s".format(
                    train_epoch, step, overall_auc, log_every / timer.tick(False)
                )
            )

    overall_auc, _, _, _ = new_metrics.calculate_overall_auc(np.array(pred_list), np.array(y_list))
    user_auc = new_metrics.calculate_user_auc(buffer)
    overall_logloss = new_metrics.calculate_overall_logloss(np.array(pred_list), np.array(y_list))
    user_ndcg5, user_hr5 = new_metrics.calculate_user_ndcg_hr(5, buffer)
    user_ndcg10, user_hr10 = new_metrics.calculate_user_ndcg_hr(10, buffer)
    user_ndcg20, user_hr20 = new_metrics.calculate_user_ndcg_hr(20, buffer)

    print("train_epoch={}, train_step={}, overall_auc={:5f}, user_auc={:5f}, overall_logloss={:5f}, "
          "user_ndcg5={:5f}, user_hr5={:5f}, user_ndcg10={:5f}, user_hr10={:5f}, user_ndcg20={:5f}, user_hr20={:5f}".
          format(train_epoch, train_step, overall_auc, user_auc, overall_logloss,
                 user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20))
    with open(os.path.join(args.checkpoint_dir, "log_ood.txt"), "a") as writer:
        print("train_epoch={}, train_step={}, overall_auc={:5f}, user_auc={:5f}, overall_logloss={:5f}, "
              "user_ndcg5={:5f}, user_hr5={:5f}, user_ndcg10={:5f}, user_hr10={:5f}, user_ndcg20={:5f}, user_hr20={:5f}".
              format(train_epoch, train_step, overall_auc, user_auc, overall_logloss,
                     user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20), file=writer)

    return overall_auc, user_auc, overall_logloss, user_ndcg5, user_hr5, user_ndcg10, user_hr10, user_ndcg20, user_hr20


def train(train_dataset, model_obj, device, args, pred_dataloader):
    criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(
        model_obj.parameters(),
        lr=float(args.learning_rate)
    )
    model_obj.train()
    model_obj.to(device)

    print(model_obj)

    logger.info("Start training...")
    timer = Timer()
    log_every = 200
    max_step = 0
    best_auc = 0
    best_auc_ckpt_path = os.path.join(args.checkpoint_dir, "best_auc" + ".pkl")
    for epoch in range(1, args.max_epoch + 1):
        for step, batch_data in enumerate(train_dataset, 1):
            logits = model_obj({
                key: value.to(device)
                for key, value in batch_data.items()
                if key not in {consts.FIELD_USER_ID, consts.FIELD_LABEL, consts.FIELD_TRIGGER_SEQUENCE}
            })

            loss = criterion(logits, batch_data[consts.FIELD_LABEL].view(-1, 1).to(device))
            pred, y = torch.sigmoid(logits), batch_data[consts.FIELD_LABEL].view(-1, 1)
            overall_auc, _, _, _ = new_metrics.calculate_overall_auc(np.array(pred.detach().cpu()), np.array(y))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if step % log_every == 0:
                logger.info(
                    "epoch={}, step={}, loss={:5f}, overall_auc={:5f}, speed={:2f} steps/s".format(
                        epoch, step, float(loss.item()), overall_auc, log_every / timer.tick(False)
                    )
                )
            max_step = step

        pred_overall_auc, pred_user_auc, pred_overall_logloss, pred_user_ndcg5, pred_user_hr5, \
        pred_user_ndcg10, pred_user_hr10, pred_user_ndcg20, pred_user_hr20 = predict(
            predict_dataset=pred_dataloader,
            model_obj=model_obj,
            device=device,
            args=args,
            train_epoch=epoch,
            train_step=epoch * max_step,
        )
        logger.info("dump checkpoint for epoch {}".format(epoch))
        model_obj.train()
        if pred_overall_auc > best_auc:
            best_auc = pred_overall_auc
            torch.save(model_obj, best_auc_ckpt_path)


def main_worker(_):
    args = parse_args()
    ap.print_arguments(args)

    args.checkpoint_dir = os.path.join(args.checkpoint_dir, "base")
    if not os.path.exists(args.checkpoint_dir):
        os.makedirs(args.checkpoint_dir)
    model_meta = model.get_model_meta(args.model)  # type: model.ModelMeta

    model_conf, raw_model_conf = ap.parse_arch_config_from_args(model_meta, args)  # type: dict

    # Construct model
    model_obj = model_meta.model_builder(model_conf=model_conf)  # type: torch.nn.module
    print("=" * 100)
    for name, parms in model_obj.named_parameters():
        print(name)
    print("=" * 100)
    device = env.get_device()
    worker_id = worker_count = 8
    train_file, test_file = args.dataset.split(',')

    args.num_loading_workers = 1
    # Setup up data loader
    # train_dataloader = sequence_dataloader.SequenceDataLoader(
    train_dataloader = sequence_dataloader.MetaSequenceDataLoader(
        table_name=train_file,
        slice_id=0,
        slice_count=args.num_loading_workers,
        is_train=True
    )
    train_dataloader = torch.utils.data.DataLoader(
        dataset=train_dataloader,
        batch_size=args.batch_size,
        num_workers=args.num_loading_workers,
        pin_memory=True,
        collate_fn=train_dataloader.batchify,
        drop_last=False,
        # shuffle=True
    )

    # Setup up data loader
    pred_dataloader = sequence_dataloader.MetaSequenceDataLoader(
        table_name=test_file,
        slice_id=args.num_loading_workers * worker_id,
        slice_count=args.num_loading_workers * worker_count,
        is_train=False
    )
    pred_dataloader = torch.utils.data.DataLoader(
        dataset=pred_dataloader,
        batch_size=args.batch_size,
        num_workers=args.num_loading_workers,
        pin_memory=True,
        collate_fn=pred_dataloader.batchify,
        drop_last=False,
    )

    # Setup training
    train(
        train_dataset=train_dataloader,
        model_obj=model_obj,
        device=device,
        args=args,
        pred_dataloader=pred_dataloader
    )


if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    mp.spawn(main_worker, nprocs=1)


main/multi_metric_vanilla_train.py
HelloZicky/DUET
from . import model_fn
from . import config

__all__ = []

model/meta_sasrec/__init__.py
HelloZicky/DUET
from ..model_meta import MetaType, model


class ModelConfig(object):
    def __init__(self):
        self.id_dimension = 8
        self.id_vocab = 500
        self.mlp_layers = 2
        self.nhead = 4

    @staticmethod
    @model("meta_sasrec", MetaType.ConfigParser)
    def parse(json_obj):
        conf = ModelConfig()
        conf.id_dimension = json_obj.get("id_dimension")
        conf.id_vocab = json_obj.get("id_vocab")
        conf.mlp_layers = json_obj.get("mlp_layers")
        conf.nhead = json_obj.get("nhead")

        return conf

model/meta_sasrec/config.py
HelloZicky/DUET
import torch.nn
import torch.nn as nn

from module import encoder, common, initializer
from util import consts
from . import config
from ..model_meta import MetaType, model

import logging
logger = logging.getLogger(__name__)


@model("meta_sasrec", MetaType.ModelBuilder)
class SasRec(nn.Module):
    def __init__(self, model_conf):
        super(SasRec, self).__init__()

        assert isinstance(model_conf, config.ModelConfig)

        self._position_embedding = encoder.IDEncoder(
            model_conf.id_vocab,
            model_conf.id_dimension
        )

        self._id_encoder = encoder.IDEncoder(
            model_conf.id_vocab,
            model_conf.id_dimension
        )

        self._target_trans = common.StackedDense(
            model_conf.id_dimension, [model_conf.id_dimension] * 2, [torch.nn.Tanh, None]
        )
        self._seq_trans = common.StackedDense(
            model_conf.id_dimension, [model_conf.id_dimension], [torch.nn.Tanh]
        )

        self._transformer = nn.TransformerEncoderLayer(
            d_model=model_conf.id_dimension,
            nhead=model_conf.nhead,
            dim_feedforward=4*model_conf.id_dimension,
            dropout=0
        )

        initializer.default_weight_init(self._transformer.self_attn.in_proj_weight)
        initializer.default_weight_init(self._transformer.self_attn.out_proj.weight)
        initializer.default_bias_init(self._transformer.self_attn.in_proj_bias)
        initializer.default_bias_init(self._transformer.self_attn.out_proj.bias)

        initializer.default_weight_init(self._transformer.linear1.weight)
        initializer.default_bias_init(self._transformer.linear1.bias)
        initializer.default_weight_init(self._transformer.linear2.weight)
        initializer.default_bias_init(self._transformer.linear2.bias)

        self._meta_classifier_param_list = common.HyperNetwork_FC(
            model_conf.id_dimension,
            [model_conf.id_dimension] * model_conf.mlp_layers,
            ([torch.nn.Tanh] * (model_conf.mlp_layers - 1)) + [None],
            batch=True,
            model_conf=model_conf
        )

    def forward(self, features, fig1=False):
        # Encode target item
        trigger_embed = self._id_encoder(features[consts.FIELD_TRIGGER_SEQUENCE])
        trigger_embed = self._seq_trans(trigger_embed)

        # B * D
        target_embed = self._id_encoder(features[consts.FIELD_TARGET_ID])
        target_embed = self._target_trans(target_embed)

        # Encode user historical behaviors
        with torch.no_grad():
            click_seq = features[consts.FIELD_CLK_SEQUENCE]
            batch_size = int(click_seq.shape[0])
            # B * L
            positions = torch.arange(0, int(click_seq.shape[1]), dtype=torch.int32).to(click_seq.device)
            positions = torch.tile(positions.unsqueeze(0), [batch_size, 1])
            mask = torch.not_equal(click_seq, 0)
            # B
            seq_length = torch.maximum(torch.sum(mask.to(torch.int32), dim=1) - 1, torch.Tensor([0]).to(device=mask.device))
            seq_length = seq_length.to(torch.long)

            trigger_mask = torch.not_equal(features[consts.FIELD_TRIGGER_SEQUENCE], 0).to(dtype=torch.int32)
            # B
            trigger_seq_length = torch.maximum(torch.sum(trigger_mask, dim=1) - 1,
                                               torch.Tensor([0]).to(device=trigger_mask.device))
            trigger_seq_length = trigger_seq_length.to(torch.long)

        # B * L * D
        hist_embed = self._id_encoder(click_seq)
        hist_pos_embed = self._position_embedding(positions)
        hist_embed = self._seq_trans(hist_embed + hist_pos_embed)

        atten_embed = self._transformer(
            torch.swapaxes(hist_embed, 0, 1)
        )
        user_state = torch.swapaxes(atten_embed, 0, 1)[range(batch_size), seq_length, :]

        user_embedding = self._meta_classifier_param_list(user_state, hist_embed,
                                         user_state.size()[0],
                                         seq_length)
        if fig1:
            user_embedding1 = self._meta_classifier_param_list(user_state, trigger_embed,
                                                              user_state.size()[0],
                                                              trigger_seq_length)
            output1 = torch.sum(user_embedding1 * target_embed, dim=1, keepdim=True)
            return output, mis_rec_pred, request_num, total_num, output1
        return torch.sum(user_embedding * target_embed, dim=1, keepdim=True)

model/meta_sasrec/model_fn.py
HelloZicky/DUET
# -*- coding: utf-8 -*-
from collections import defaultdict


__all__ = ["ModelMeta", "get_model_meta", "model"]


class ModelMeta(object):
    def __init__(self, config_parser=None, data_loader=None, model_builder=None):
        self._config_parser = config_parser
        self._data_loader = data_loader
        self._model_builder = model_builder

    @property
    def arch_config_parser(self):
        return self._config_parser

    def set_arch_config_parser(self, parser):
        self._check(self._config_parser, "Config parser has been set")
        self._config_parser = parser

    @property
    def data_loader_builder(self):
        return self._data_loader

    def set_data_loader_builder(self, loader):
        self._check(self._data_loader, "Data loader builder has been set")
        self._data_loader = loader

    @property
    def model_builder(self):
        return self._model_builder

    def set_model_builder(self, model_builder):
        self._check(self._model_builder, "Model builder has been set")
        self._model_builder = model_builder

    def _check(self, value, message):
        if value is not None:
            raise ValueError(message)

    def __setitem__(self, k, v):
        self.k = v


class MetaType(object):
    ConfigParser = ModelMeta.set_arch_config_parser
    ModelBuilder = ModelMeta.set_model_builder


class _ModelMetaRegister(object):
    def __init__(self):
        self._register_map = defaultdict(ModelMeta)

    def get(self, name):
        return self._register_map.get(name)

    def __call__(self, name, setter):
        model_meta = self._register_map[name]

        def _executor(func):
            setter(model_meta, func)
            return func

        return _executor


model = _ModelMetaRegister()
get_model_meta = model.get

model/model_meta.py
HelloZicky/DUET
from . import model_fn
from . import config

__all__ = []

model/sasrec/__init__.py
HelloZicky/DUET
from ..model_meta import MetaType, model


class ModelConfig(object):
    def __init__(self):
        self.id_dimension = 8
        self.id_vocab = 500
        self.mlp_layers = 2
        self.nhead = 4

    @staticmethod
    @model("sasrec", MetaType.ConfigParser)
    def parse(json_obj):
        conf = ModelConfig()
        conf.id_dimension = json_obj.get("id_dimension")
        conf.id_vocab = json_obj.get("id_vocab")
        conf.mlp_layers = json_obj.get("mlp_layers")
        conf.nhead = json_obj.get("nhead")

        return conf

model/sasrec/config.py
HelloZicky/DUET
import torch.nn
import torch.nn as nn

from module import encoder, common, initializer
from util import consts
from . import config
from ..model_meta import MetaType, model

import logging
logger = logging.getLogger(__name__)


@model("sasrec", MetaType.ModelBuilder)
class SasRec(nn.Module):
    def __init__(self, model_conf):
        super(SasRec, self).__init__()

        assert isinstance(model_conf, config.ModelConfig)

        self._position_embedding = encoder.IDEncoder(
            model_conf.id_vocab,
            model_conf.id_dimension
        )

        self._id_encoder = encoder.IDEncoder(
            model_conf.id_vocab,
            model_conf.id_dimension
        )

        self._target_trans = common.StackedDense(
            model_conf.id_dimension, [model_conf.id_dimension] * 2, [torch.nn.Tanh, None]
        )
        self._seq_trans = common.StackedDense(
            model_conf.id_dimension, [model_conf.id_dimension], [torch.nn.Tanh]
        )

        self._transformer = nn.TransformerEncoderLayer(
            d_model=model_conf.id_dimension,
            nhead=model_conf.nhead,
            dim_feedforward=4*model_conf.id_dimension,
            dropout=0
        )

        initializer.default_weight_init(self._transformer.self_attn.in_proj_weight)
        initializer.default_weight_init(self._transformer.self_attn.out_proj.weight)
        initializer.default_bias_init(self._transformer.self_attn.in_proj_bias)
        initializer.default_bias_init(self._transformer.self_attn.out_proj.bias)

        initializer.default_weight_init(self._transformer.linear1.weight)
        initializer.default_bias_init(self._transformer.linear1.bias)
        initializer.default_weight_init(self._transformer.linear2.weight)
        initializer.default_bias_init(self._transformer.linear2.bias)

        self._classifier = common.StackedDense(
            model_conf.id_dimension,
            [model_conf.id_dimension] * model_conf.mlp_layers,
            ([torch.nn.Tanh] * (model_conf.mlp_layers - 1)) + [None]
        )

    def forward(self, features):
        # Encode target item
        # B * D
        target_embed = self._id_encoder(features[consts.FIELD_TARGET_ID])
        target_embed = self._target_trans(target_embed)

        # Encode user historical behaviors
        with torch.no_grad():
            click_seq = features[consts.FIELD_CLK_SEQUENCE]
            batch_size = int(click_seq.shape[0])
            # B * L
            positions = torch.arange(0, int(click_seq.shape[1]), dtype=torch.int32).to(click_seq.device)
            positions = torch.tile(positions.unsqueeze(0), [batch_size, 1])
            mask = torch.not_equal(click_seq, 0)
            # B
            seq_length = torch.maximum(torch.sum(mask.to(torch.int32), dim=1) - 1, torch.Tensor([0]).to(device=mask.device))
            seq_length = seq_length.to(torch.long)

        # B * L * D
        hist_embed = self._id_encoder(click_seq)
        hist_pos_embed = self._position_embedding(positions)
        hist_embed = self._seq_trans(hist_embed + hist_pos_embed)

        atten_embed = self._transformer(
            torch.swapaxes(hist_embed, 0, 1)
        )
        user_state = torch.swapaxes(atten_embed, 0, 1)[range(batch_size), seq_length, :]

        user_embedding = self._classifier(user_state)

        return torch.sum(user_embedding * target_embed, dim=1, keepdim=True)

model/sasrec/model_fn.py
HelloZicky/DUET
import torch
from torch import nn

from . import common


class TargetAttention(nn.Module):
    def __init__(self, key_dimension, value_dimension):
        super(TargetAttention, self).__init__()
        self._target_key_transform = common.Linear(key_dimension, key_dimension, bias=False)
        self._item_key_transform = common.Linear(key_dimension, key_dimension, bias=False)
        self._value_transform = common.Linear(value_dimension, value_dimension, bias=False)
        self._scaler = key_dimension ** 0.5

    def __setitem__(self, k, v):
        self.k = v

    def forward(self, target_key, item_keys, item_values, mask):
        """
        :param target_key: B * D
        :param item_keys: B * L * D
        :param item_values: B * L * D
        :param mask: B * L
        :return:
        """
        assert item_keys.shape[1] == item_values.shape[1]
        assert target_key.shape[-1] == item_keys.shape[-1]
        assert target_key.shape[0] == item_keys.shape[0] == item_values.shape[0]

        target_key = self._target_key_transform(target_key)[:, None, :]
        item_keys = self._item_key_transform(item_keys)
        item_values = self._value_transform(item_values)

        atten_weights = torch.sum(target_key * item_keys, dim=-1, keepdim=True) / self._scaler
        if mask is not None:
            atten_weights += -1e8 * (1 - mask[:, :, None])
            
        atten_weights = torch.softmax(atten_weights, dim=1)
        return torch.sum(atten_weights * item_values, dim=1)


if __name__ == '__main__':
    target_embed = torch.randn(16, 8)
    item_keys = torch.randn(16, 10, 8)
    item_values = torch.randn(16, 10, 23)

    m = TargetAttention(
        key_dimension=8,
        value_dimension=23,
        value_out_dimension=32
    )
    mask = torch.cat([torch.ones([16, 4]), torch.zeros([16, 6])], dim=1)
    data = m(target_embed, item_keys, item_values, mask)
    print(data.shape)
module/attention.py
HelloZicky/DUET
"""
Common modules
"""
import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
import torch.nn.functional as F

from . import initializer


class StackedDense(torch.nn.Module):
    def __init__(self, in_dimension, units, activation_fns):
        super(StackedDense, self).__init__()

        modules = []
        units = [in_dimension] + list(units)
        for i in range(1, len(units)):
            linear = torch.nn.Linear(units[i - 1], units[i], bias=True)
            initializer.default_weight_init(linear.weight)
            initializer.default_bias_init(linear.bias)
            modules.append(linear)

            if activation_fns[i - 1] is not None:
                modules.append(activation_fns[i - 1]())

        self.net = torch.nn.Sequential(*modules)

    def __setitem__(self, k, v):
        self.k = v

    def forward(self, x):
        return self.net(x)


class Linear(torch.nn.Module):
    def __init__(self, in_dimension, out_dimension, bias):
        super(Linear, self).__init__()
        self.net = torch.nn.Linear(in_dimension, out_dimension, bias)
        initializer.default_weight_init(self.net.weight)
        if bias:
            initializer.default_weight_init(self.net.bias)

    def __setitem__(self, k, v):
        self.k = v

    def forward(self, x):
        return self.net(x)


class HyperNetwork_FC(nn.Module):
    def __init__(self, in_dimension, units, activation_fns, batch=True, trigger_sequence_len=30,
                 model_conf=None, expand=False):
        super(HyperNetwork_FC, self).__init__()

        self.modules = []
        units = [in_dimension] + list(units)
        self.activation_fns = activation_fns
        self.batch = batch
        self.units = units
        self.w1 = []
        self.w2 = []
        self.b1 = []
        self.b2 = []
        self.output_size = []
        self._gru_cell = torch.nn.GRU(
            model_conf.id_dimension,
            model_conf.id_dimension,
            batch_first=True
        )
        self._mlp_trans = StackedDense(
            model_conf.id_dimension,
            [model_conf.id_dimension] * model_conf.mlp_layers,
            ([torch.nn.Tanh] * (model_conf.mlp_layers - 1)) + [None]
        )
        initializer.default_weight_init(self._gru_cell.weight_hh_l0)
        initializer.default_weight_init(self._gru_cell.weight_ih_l0)
        initializer.default_bias_init(self._gru_cell.bias_ih_l0)
        initializer.default_bias_init(self._gru_cell.bias_hh_l0)
        self.expand = expand
        for i in range(1, len(units)):
            if i == 1 and self.expand:
                input_size = units[i - 1] * 2
            else:
                input_size = units[i - 1]

            output_size = units[i]

            self.w1.append(Parameter(torch.fmod(torch.randn(in_dimension, input_size * output_size).cuda(), 2)))
            self.b1.append(Parameter(torch.fmod(torch.randn(input_size * output_size).cuda(), 2)))

            self.w2.append(Parameter(torch.fmod(torch.randn(in_dimension, output_size).cuda(), 2)))
            self.b2.append(Parameter(torch.fmod(torch.randn(output_size).cuda(), 2)))

            if activation_fns[i - 1] is not None:
                self.modules.append(activation_fns[i - 1]())

    def forward(self, x, z, sample_num=32, trigger_seq_length=30):
        units = self.units

        user_state, _ = self._gru_cell(z)
        user_state = user_state[range(user_state.shape[0]), trigger_seq_length, :]
        z = self._mlp_trans(user_state)

        for i in range(1, len(units)):
            index = i - 1
            if i == 1 and self.expand:
                input_size = units[i - 1] * 2
            else:
                input_size = units[i - 1]

            weight = torch.matmul(z, self.w1[index]) + self.b1[index]

            output_size = units[i]

            if not self.batch:
                weight = weight.view(input_size, output_size)
            else:
                weight = weight.view(sample_num, input_size, output_size)
            bias = torch.matmul(z, self.w2[index]) + self.b2[index]

            x = torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
            x = self.modules[index](x) if index < len(self.modules) else x

        return x


module/common.py
HelloZicky/DUET
"""
Basic feature encoders
"""
import torch

from . import initializer


class BaseEncoder(torch.nn.Module):
    def __init__(self):
        super(BaseEncoder, self).__init__()

    def in_dimension(self):
        raise NotImplementedError

    def out_dimension(self):
        raise NotImplementedError


class DenseEncoder(BaseEncoder):
    def __init__(self, in_dimension, out_dimension, activation=torch.nn.Tanh):
        super(DenseEncoder, self).__init__()
        self._in_dimension = in_dimension
        self._out_dimension = out_dimension
        self._fully_connect = torch.nn.Sequential(
            torch.nn.Linear(self._in_dimension, self._out_dimension),
            activation()
        )

    def __setitem__(self, k, v):
        self.k = v

    def forward(self, x):
        return self._fully_connect(x)

    def out_dimension(self):
        return self._out_dimension

    def in_dimension(self):
        return self._in_dimension


class IDEncoder(BaseEncoder):
    def __init__(self, vocab_size, out_dimension):
        super(IDEncoder, self).__init__()
        self._embedding_matrix = torch.nn.Embedding(
            num_embeddings=vocab_size,
            embedding_dim=out_dimension
        )
        initializer.default_weight_init(self._embedding_matrix.weight)

    def __setitem__(self, k, v):
        self.k = v
        
    def forward(self, x):
        return self._embedding_matrix(x)

    def out_dimension(self):
        return self._embedding_matrix.embedding_dim

module/encoder.py
HelloZicky/DUET
import torch


# trunk model init
def default_weight_init(tensor):
    torch.nn.init.xavier_uniform(tensor)


def default_bias_init(tensor):
    torch.nn.init.constant_(tensor, 0)


# lite plugin model init
def default_lite_plugin_init(layer):
    torch.nn.init.xavier_uniform(layer.weight, gain=0.001)
    # torch.nn.init.constant_(layer.weight, 0)
    torch.nn.init.constant_(layer.bias, 0)


# naive plugin model init
def default_naive_plugin_init(layer):
    torch.nn.init.constant_(layer.weight, 0)
    torch.nn.init.constant_(layer.bias, 0)


# Initializers will be optimized in the future
#
# def weight_init_normal(m):
#     if isinstance(m, nn.Linear) or isinstance(m, nn.Embedding):
#         torch.nn.init.normal_(m.weight, mean=0, std=0.01)
#     if isinstance(m, nn.Linear):
#         torch.nn.init.normal_(m.bias, mean=0, std=0.01)
#
#
# def weight_init_constant(m):
#     if isinstance(m, nn.Linear) or isinstance(m, nn.Embedding):
#         torch.nn.init.constant_(m.weight, 0)
#     if isinstance(m, nn.Linear):
#         torch.nn.init.constant_(m.bias, 0)


if __name__ == '__main__':
    # model.apply(weight_init_normal)
    dimension = 10
    plugin_layer = torch.nn.Linear(dimension, dimension // 2, True)
    print("-" * 50)
    print("original")
    print("plugin_layer.weight", plugin_layer.weight)
    print("plugin_layer.bias", plugin_layer.bias)
    default_weight_init(plugin_layer.weight)
    default_bias_init(plugin_layer.bias)
    print("-" * 50)
    print("trunk_init")
    print("plugin_layer.weight", plugin_layer.weight)
    print("plugin_layer.bias", plugin_layer.bias)
    default_lite_plugin_init(plugin_layer)
    print("-" * 50)
    print("lite_plugin_init")
    print("plugin_layer.weight", plugin_layer.weight)
    print("plugin_layer.bias", plugin_layer.bias)
    default_naive_plugin_init(plugin_layer)
    print("-" * 50)
    print("naive_plugin_init")
    print("plugin_layer.weight", plugin_layer.weight)
    print("plugin_layer.bias", plugin_layer.bias)
module/initializer.py
HelloZicky/DUET
# -*- coding: utf-8 -*-
"""
Argument parsing and dumpling
"""
import copy
import json
import base64
import os.path
from argparse import Namespace

_ARCH_CONFIG_FILE_NAME = "arch_conf.json"
_TRAIN_ARGS_FILE_NAME = "train_args.json"


def parse_arch_config_from_args(model, args):
    """
    Read or parse arch config
    :param model:
    :param args:
    :return:
    """
    if args.arch_config is not None:
        with open(args.arch_config) as jsonfile:
            raw_arch_config = json.load(jsonfile)
    elif args.arch_config_path is not None:
        with open(args.arch_config_path, "rt") as reader:
            raw_arch_config = json.load(reader)
    else:
        raise KeyError("Model configuration not found")

    return model.arch_config_parser(raw_arch_config), raw_arch_config



def parse_arch_config_from_args_get_profile(model_meta, arch_config, bucket):
    """
    Read or parse arch config
    :param model_meta:
    :param args:
    :return:
    """
    print(arch_config)
    # raw_arch_config = json.loads(arch_config)
    f = open(arch_config, encoding="utf-8")
    raw_arch_config = json.load(f)

    return model_meta.arch_config_parser(raw_arch_config), raw_arch_config


def dump_config(checkpoint_dir, file_name, config_obj):
    """
    Dump configurations to OSS
    :param checkpoint_dir:
    :param file_name:
    :param config_obj:
    :return:
    """
    print(config_obj, file=open(os.path.join(checkpoint_dir, file_name), "w+"))


# def dump_model_config(checkpoint_dir, raw_model_arch, bucket):
def dump_model_config(checkpoint_dir, raw_model_arch):
    """
    Dump model configurations to OSS
    :param args: Namespace object, parsed from command-line arguments
    :param raw_model_arch:
    :return:
    """
    dump_config(checkpoint_dir, _ARCH_CONFIG_FILE_NAME, raw_model_arch)


# def dump_train_arguments(checkpoint_dir, args, bucket):
def dump_train_arguments(checkpoint_dir, args):
    args_dict = copy.copy(args.__dict__)
    args_dict.pop("arch_config")
    dump_config(checkpoint_dir, _TRAIN_ARGS_FILE_NAME, args_dict)


def print_arguments(args):
    print("\nMain arguments:")
    for k, v in args.__dict__.items():
        print("{}={}".format(k, v))

util/args_processing.py
HelloZicky/DUET

FIELD_USER_ID = "user_id"
FIELD_TARGET_ID = "target_id"
# FIELD_REQ_SEQUENCE = "req_sequence"
FIELD_CLK_SEQUENCE = "clk_sequence"
FIELD_TRIGGER_SEQUENCE = "trigger_sequence"
FIELD_LABEL = "label"
FIELD_BATCH = "batch"

STATE_MODEL = "model"
STATE_OPTIMIZER = "optimizer"
STATE_STEP = "step"

FILE_CHECKPOINT_PREFIX = "ckpt-"

ENVIRON_WORLD_SIZE = "WORLD_SIZE"
ENVIRON_RANK = "RANK"

util/consts.py
HelloZicky/DUET
import logging
import os

import torch

from . import consts

logger = logging.getLogger(__name__)


def get_device():
    try:
        logger.info(torch.cuda.get_device_name(torch.cuda.current_device()))
        device = torch.device("cuda")
    except:
        device = torch.device("cpu")

    return device


def get_cluster_info():
    return int(os.environ[consts.ENVIRON_RANK]), int(os.environ[consts.ENVIRON_WORLD_SIZE])
util/env.py
HelloZicky/DUET
from sklearn import metrics
import heapq
import numpy as np
import math
import itertools
import torch


def calculate_user_ndcg_hr(n=5, *buffer):
    user_num = 0
    ndcg_ = 0
    hr_ = 0
    for user_id, actions in itertools.groupby(buffer[0], key=lambda x: x[0]):
        actions = sorted(actions, key=lambda x: x[1], reverse=True)
        top_items = np.array(actions)[:n, :]
        num_postive = int(sum(np.array(actions)[:, 2]))
        if not 0 < num_postive:
            continue
        dcg = 0
        idcg = 0
        for i, (user_id, score, label) in enumerate(top_items):
            if label == 1:
                dcg += math.log(2) / math.log(i + 2)
            if i < num_postive:
                idcg += math.log(2) / math.log(i + 2)

        ndcg_ += dcg / idcg
        hr_ += 1 if any(item[2] for item in top_items) else 0
        user_num += 1

    ndcg = ndcg_ / user_num
    hr = hr_ / user_num
    return ndcg, hr


def calculate_overall_logloss(*buffer):
    prob, y = buffer
    logloss = float(metrics.log_loss(np.array(y), prob))

    return logloss


def calculate_overall_auc(*buffer):
    prob, y = buffer
    fpr, tpr, thresholds = [], [], []
    auc = float(metrics.roc_auc_score(np.array(y), prob))

    return auc, fpr, tpr, thresholds


def calculate_user_auc(*buffer):
    user_num = 0
    auc_ = 0
    for user_id, actions in itertools.groupby(buffer[0], key=lambda x: x[0]):
        actions = list(actions)
        prob = np.array(actions)[:, 1]
        y = np.array(actions)[:, 2]
        if not 0 < np.sum(y) < len(y):
            continue
        auc_ += float(metrics.roc_auc_score(np.array(y), prob))
        user_num += 1
    auc = auc_ / user_num
    return auc


def calculate_user_auc_misrec(*buffer):
    user_num = 0
    auc_ = 0
    fig1_list = []
    for user_id, actions in itertools.groupby(buffer[0], key=lambda x: x[0]):
        actions = list(actions)
        prob = np.array(actions)[:, 1]
        prob_trigger = np.array(actions)[:, 2]
        y = np.array(actions)[:, 3]
        mis_rec = np.mean(np.array(actions)[:, 4])
        auc_item = float(metrics.roc_auc_score(np.array(y), prob))
        auc_item_trigger = float(metrics.roc_auc_score(np.array(y), prob_trigger))
        auc_ += auc_item
        fig1_list.append([auc_item, auc_item_trigger, mis_rec])
        user_num += 1
    auc = auc_ / user_num
    return auc, fig1_list


def calculate_overall_acc(*buffer):
    prob, y = buffer
    acc = metrics.accuracy_score(y, prob)

    return acc


def calculate_overall_roc(*buffer):
    prob, y = buffer
    fpr, tpr, thresholds = metrics.roc_curve(np.array(y), prob, pos_label=1)

    return fpr, tpr


def calculate_overall_recall(*buffer):
    prob, y = buffer
    recall = metrics.recall_score(y, prob, average="macro")
    return recall


def calculate_overall_precision(*buffer):
    prob, y = buffer
    precision = metrics.precision_score(y, prob, average="macro")
    return precision


if __name__ == '__main__':
    records = [
        [3, 0.151, 0],
        [3, 0.152, 0],
        [3, 0.103, 0],
        [3, 0.174, 1],
        [3, 0.135, 0],
        [3, 0.126, 0],
        [1, 0.151, 1],
        [1, 0.152, 0],
        [1, 0.103, 0],
        [1, 0.174, 0],
        [1, 0.135, 0],
        [2, 0.151, 0],
        [2, 0.152, 0],
        [2, 0.103, 1],
        [2, 0.174, 0],
        [2, 0.135, 0],
        [2, 0.126, 0],
    ]
    # ndcg, hr = calculate_user_ndcg_hr(np.array(records)[:, 0], np.array(records)[:, 1], np.array(records)[:, 2])
    # ndcg, hr = calculate_user_ndcg_hr(np.hstack(np.array(records)[:, 0], np.array(records)[:, 1], np.array(records)[:, 2]))
    # ndcg, hr = calculate_user_ndcg_hr(torch.Tensor(np.array(records)[:, 0]).view(-1, 1),
    #                                   torch.Tensor(np.array(records)[:, 1]).view(-1, 1),
    #                                   torch.Tensor(np.array(records)[:, 2]).view(-1, 1))
    # user_id_list = torch.Tensor([3, 3, 3, 3, 3, 3, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2,]).view(-1, 1)
    # score_list = torch.Tensor([0.15, 0.15, 0.10, 0.17, 0.13, 0.12, 0.15, 0.15, 0.10, 0.17, 0.13, 0.15, 0.15, 0.10, 0.17, 0.13, 0.12,]).view(-1, 1)
    # label_list = torch.Tensor([0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,]).view(-1, 1)
    # print(torch.Tensor(np.array(records)[:, 0]).view(-1, 1))
    # ndcg, hr = calculate_user_ndcg_hr(torch.cat(torch.Tensor(np.array(records)[:, 0]).view(-1, 1),
    #                                   torch.Tensor(np.array(records)[:, 1]).view(-1, 1),
    #                                   torch.Tensor(np.array(records)[:, 2]).view(-1, 1), dim=1))
    # ndcg, hr = calculate_user_ndcg_hr(user_id_list, score_list, label_list)
    ndcg, hr = calculate_user_ndcg_hr(records)
    print(ndcg, hr)

util/new_metrics.py
HelloZicky/DUET
import time


class Timer:
    def __init__(self):
        self._start = time.time()
        self._last = self._start

    def tick(self, round=True):
        ret_value = time.time() - self._last
        self._last += ret_value
        if round:
            ret_value = int(ret_value)
        return ret_value

    def total_time(self):
        return time.time() - self._start

util/timer.py
HelloZicky/DUET
from collections import defaultdict, deque
import datetime
import time
import torch
import torch.distributed as dist
import logging
from io import BytesIO
logger = logging.getLogger(__name__)


def setup_seed(seed):
    import numpy as np
    import random
    from torch.backends import cudnn
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.deterministic = True
    cudnn.benchmark = False

    
def is_dist_avail_and_initialized():
    if not dist.is_available():
        return False
    if not dist.is_initialized():
        return False
    return True


class SmoothedValue(object):
    """Track a series of values and provide access to smoothed values over a
    window or the global series average.
    """

    def __init__(self, window_size=20, fmt=None):
        if fmt is None:
            fmt = "{median:.4f} ({global_avg:.4f})"
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    def synchronize_between_processes(self):
        """
        Warning: does not synchronize the deque!
        """
        if not is_dist_avail_and_initialized():
            return
        print("count / total: {} / {}".format(self.count, self.total))
        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
        dist.barrier()
        dist.all_reduce(t)
        t = t.tolist()
        self.count = int(t[0])
        self.total = t[1]
        print("count / total: {} / {}".format(self.count, self.total))

    @property
    def median(self):
        d = torch.tensor(list(self.deque))
        return d.median().item()

    @property
    def avg(self):
        d = torch.tensor(list(self.deque), dtype=torch.float32)
        return d.mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deque)

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value)


class MetricLogger(object):
    def __init__(self, delimiter="\t"):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    def update(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            assert isinstance(v, (float, int))
            self.meters[k].update(v)

    def __getattr__(self, attr):
        if attr in self.meters:
            return self.meters[attr]
        if attr in self.__dict__:
            return self.__dict__[attr]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, attr))

    def __str__(self):
        loss_str = []
        for name, meter in self.meters.items():
            loss_str.append(
                "{}: {}".format(name, str(meter))
            )
        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):
        # for meter in self.meters.values():
        for name, meter in self.meters.items():
            print("Measure: {}".format(name))
            meter.synchronize_between_processes()

    def add_meter(self, name, meter):
        self.meters[name] = meter

    def log_every(self, iterable, print_freq, total_iter, header=None):
        i = 0
        if not header:
            header = ''
        start_time = time.time()
        end = time.time()
        iter_time = SmoothedValue(fmt='{avg:.4f}')
        data_time = SmoothedValue(fmt='{avg:.4f}')
        space_fmt = ':' + str(len(str(total_iter))) + 'd'
        if torch.cuda.is_available():
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}',
                'max mem: {memory:.0f}'
            ])
        else:
            log_msg = self.delimiter.join([
                header,
                '[{0' + space_fmt + '}/{1}]',
                'eta: {eta}',
                '{meters}',
                'time: {time}',
                'data: {data}'
            ])
        MB = 1024.0 * 1024.0
        for obj in iterable:
            data_time.update(time.time() - end)
            yield obj, i
            iter_time.update(time.time() - end)
            if i % print_freq == 0:
                eta_seconds = iter_time.global_avg * (total_iter - i)
                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
                if torch.cuda.is_available():
                    logger.info(log_msg.format(
                        i, total_iter, eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time),
                        memory=torch.cuda.max_memory_allocated() / MB))
                else:
                    logger.info(log_msg.format(
                        i, total_iter, eta=eta_string,
                        meters=str(self),
                        time=str(iter_time), data=str(data_time)))
            i += 1
            end = time.time()
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        logger.info('{} Total time: {}'.format(header, total_time_str))


def ce_accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target[None])

        res = []
        for k in topk:
            correct_k = correct[:k].flatten().sum(dtype=torch.float32)
            res.append(correct_k * (100.0 / batch_size))
        return res


def bce_accuracy(output, target, threshold=0.5):
    with torch.no_grad():
        output = torch.nn.Sigmoid()(output)
        print(output)
        zeros = torch.zeros_like(output)
        ones = torch.ones_like(output)
        output = torch.where(output >= threshold, ones, zeros)
        print(output)
        pred = output * target
        # batch_size = target.size(0)
        # num_classes = target.size(1)

        TP = pred.flatten().sum(dtype=torch.float32)
        TPFP = output.flatten().sum(dtype=torch.float32)
        TPFN = target.flatten().sum(dtype=torch.float32)

        return TP * (100.0 / TPFP), TP * (100.0 / TPFN)


def kldiv_accuracy(output, target, topk=(1,)):
    """Computes the accuracy over the k top predictions for the specified values of k"""
    with torch.no_grad():
        TPFN = target.flatten().sum(dtype=torch.float32)

        res = []
        for k in topk:
            values, indices = output.topk(k, 1, True, True)
            pred = torch.zeros_like(output)
            ones = torch.ones_like(output)
            pred.scatter_(1, indices, ones)
            match_result = pred * target

            TP = match_result.flatten().sum(dtype=torch.float32)
            TPFP = pred.flatten().sum(dtype=torch.float32)
            res.append((TP, TPFP))

        return res, TPFN


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    def print_pass(*args, **kwargs):
        pass

    __builtin__.print = print
    if not is_master:
        logger.info = print_pass


def init_distributed_mode(args, gpu, ngpus_per_node, world_size, rank):
    assert (torch.cuda.is_available())
    torch.cuda.set_device(gpu)
    args.gpu = gpu
    args.dist_rank = rank * ngpus_per_node + gpu
    args.world_size = world_size
    args.rank = rank
    if world_size == 1:
        logger.info('Not using distributed mode')
        args.distributed = False
        return

    args.dist_backend = 'nccl'
    dist.init_process_group(backend=args.dist_backend, world_size=args.world_size, rank=args.dist_rank)
    args.distributed = True


util/utils.py
HelloZicky/DUET