train_ldpc.py

import argparse
import datetime
import logging
import numpy as np
import os
import sys
import torch
from tensorboardX import SummaryWriter
import time
from tqdm import tqdm

import lib
from lib.model.mpnn import factor_mpnn, FactorNN
import scipy.stats as st
import utils
from utils.types import str2bool, to_cuda


class LDPCModel(torch.nn.Module):
    def __init__(self, nfeature_dim, hop_order, nedge_type, with_residual=True, aggregator='max'):
        super(LDPCModel, self).__init__()

        self.main = FactorNN(nfeature_dim,
                             [hop_order, 96],
                             [64, 64, 64, 128, 256, 256, 128, 64, 64],
                             [nedge_type, 1],
                             2,
                             skip_link={4: 3, 5: 2, 7: 0},
                             ret_high=True,
                             aggregator=aggregator)

        self.emodel_f2v = torch.nn.Sequential(torch.nn.Conv2d(7, 64, 1),
                                              torch.nn.ReLU(inplace=True),
                                              torch.nn.Conv2d(64, nedge_type, 1))

        self.emodel_v2f = torch.nn.Sequential(torch.nn.Conv2d(7, 64, 1),
                                              torch.nn.ReLU(inplace=True),
                                              torch.nn.Conv2d(64, nedge_type, 1))

        hetype_v2f = np.ones([1, 1, 1, 96]).astype(np.float32)
        hetype_f2v = np.ones([1, 1, 96, 1]).astype(np.float32)

        hnn_idx_v2f = np.asarray(
            list(range(96)), dtype=np.int64).reshape(1, 1, 96)
        hnn_idx_f2v = np.asarray([0] * 96, dtype=np.int64).reshape(1, 96, 1)

        self.hnn_idx_v2f = torch.nn.Parameter(
            torch.from_numpy(hnn_idx_v2f).long(), requires_grad=False)
        self.hnn_idx_f2v = torch.nn.Parameter(
            torch.from_numpy(hnn_idx_f2v).long(), requires_grad=False)

        self.hetype_v2f = torch.nn.Parameter(
            torch.from_numpy(hetype_v2f), requires_grad=False)
        self.hetype_f2v = torch.nn.Parameter(
            torch.from_numpy(hetype_f2v), requires_grad=False)

        self.with_residual = with_residual

        self.nhop_regressor = torch.nn.Sequential(torch.nn.Linear(64, 128),
                                                  torch.nn.BatchNorm1d(128),
                                                  torch.nn.ReLU(),
                                                  torch.nn.Linear(128, 128),
                                                  torch.nn.ReLU(),
                                                  torch.nn.Linear(128, 1),
                                                  torch.nn.ReLU())

    def forward(self, node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f):
        etype_f2v = self.emodel_f2v(efeature_f2v)
        etype_v2f = self.emodel_v2f(efeature_v2f)

        with torch.no_grad():
            bsize = node_feature.shape[0]
            nhop_feature = node_feature[:, 0, :, :]
            nhop_feature = nhop_feature.reshape(bsize, 96, 1, 1)

        res, nhops = self.main(node_feature,
                               [hop_feature, nhop_feature],
                               [nn_idx_f2v, self.hnn_idx_f2v.repeat(
                                   bsize, 1, 1)],
                               [nn_idx_v2f, self.hnn_idx_v2f.repeat(
                                   bsize, 1, 1)],
                               [etype_f2v, self.hetype_f2v.repeat(
                                   bsize, 1, 1, 1)],
                               [etype_v2f, self.hetype_v2f.repeat(bsize, 1, 1, 1)])

        if self.with_residual:
            res = res + node_feature[:, :1, :, :]
        batch_size = res.shape[0]
        res = res.squeeze()
        if batch_size == 1:
            res = res.unsqueeze(0)

        hhop = nhops[1].squeeze()
        if bsize == 1:
            hhop = hhop.unsqueeze(0)

        snr_b_pred = self.nhop_regressor(hhop)

        return res[:, :48].contiguous(), snr_b_pred


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--n_epochs',
                        type=int,
                        default=10,
                        help="training epoches")
    parser.add_argument('--model_path',
                        type=str,
                        default='gnn',
                        help="Saved model path")
    parser.add_argument('--model_name',
                        type=str,
                        default='FactorNN',
                        help="model name")
    parser.add_argument('--use_cuda',
                        type=str2bool,
                        default=True,
                        help="Use cuda or not")

    parser.add_argument('--snr', type=int, default=None, help="snr")

    parser.add_argument('--test_path',
                        type=str,
                        default="dataset/ldpc_valid.pt",
                        help="path of the testing dataset")

    parser.add_argument('--filename', type=str, default='ldpc_data/96.3.963')
    parser.add_argument('--train', action='store_true', default=False)

    parser.add_argument('--batch_size', type=int, default=32)
    parser.add_argument('--aggregator', type=str, default='max')
    args = parser.parse_args()

    if not torch.cuda.is_available():
        args.use_cuda = False

    return args


def worker_init_fn(idx):
    t = int(time.time() * 1000.0) + idx
    seed = ((t & 0xff000000) >> 24) + ((t & 0x00ff0000) >> 8) + \
        ((t & 0x0000ff00) << 8) + ((t & 0x000000ff) << 24)
    np.random.seed(seed)
    lib.data.init_seed(int(seed % 65537))


def train(args, model,  writer, model_dir):

    train_dataset = lib.data.ContinousCodesSP(snr=args.snr)

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=8,
                                               worker_init_fn=worker_init_fn)

    optimizer = torch.optim.Adam(
        model.parameters(), lr=1e-2,  weight_decay=1e-8)

    def lr_sched(x, start=10):
        if x <= start:
            return max(1e-2, (1.0 / start) * x)
        else:
            return max(0.99 ** (x - start), 1e-6)
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lr_lambda=lambda x: lr_sched(x))
    start_epoch = 0
    gcnt = 0
    if os.path.exists(args.model_path):
        if not torch.cuda.is_available():
            ckpt = torch.load(
                args.model_path, map_location=torch.device('cpu'))
        else:
            ckpt = torch.load(args.model_path)
        model.load_state_dict(ckpt['model_state_dict'])
        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        scheduler.load_state_dict(ckpt['lr_sche'])

        start_epoch = ckpt['epoch']
        gcnt = ckpt['gcnt']

    def get_model_dict():
        return {
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'lr_sche': scheduler.state_dict(),
            'epoch': epoch,
            'gcnt': gcnt
        }

    def get_model_path():
        return os.path.join(model_dir, '{}_nn_factor_epoches_{}_snr_{}.pt'.format(args.model_name, epoch, args.snr))

    print('training started!')
    for epoch in tqdm(range(start_epoch, args.n_epochs)):
        if (epoch + 1) % 10 == 0:
            torch.save(get_model_dict(), get_model_path())

        logging.info(f'save train result to {get_model_path()}')

        loss_seq = []
        sigma_b_loss_seq = []
        acc_seq = []
        for bcnt, (node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b) in tqdm(enumerate(train_loader)):
            optimizer.zero_grad()
            if args.use_cuda:
                node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b = to_cuda(
                    node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b.float())

            if len(node_feature.shape) == 3:
                node_feature = node_feature.unsqueeze(-1)

            pred, sigma_b_pred = model(node_feature, hop_feature, nn_idx_f2v,
                                       nn_idx_v2f, efeature_f2v, efeature_v2f)

            label = label[:, :48].contiguous()
            # print(pred.shape)
            # print(label.shape)
            loss = torch.nn.functional.binary_cross_entropy_with_logits(
                pred.view(-1), label.view(-1).float())
            sigma_b_loss = torch.nn.functional.mse_loss(
                sigma_b_pred.view(-1), (torch.pow(10.0, sigma_b.float()/20)).view(-1))

            allloss = loss + 0.1 * sigma_b_loss
            allloss.backward()
            # torch.nn.utils.clip_grad_norm(parameters, 1.0)

            optimizer.step()
            loss_seq.append(loss.item())
            gcnt += 1

            pred_int = (pred > 0)
            all_correct = torch.sum(pred_int.long() == label)
            acc = all_correct.item() / np.prod(label.shape)

            acc_seq.append(acc)
            sigma_b_loss_seq.append(sigma_b_loss.item())

            if gcnt % 10 == 0:
                print('epoch = {} bcnt = {} loss = {} acc = {}'.format(
                    epoch, bcnt, np.mean(loss_seq), np.mean(acc_seq)))
                writer.add_scalar('syn_train/loss', np.mean(loss_seq), gcnt)
                writer.add_scalar('syn_train/sigma_b_loss',
                                  np.mean(sigma_b_loss_seq), gcnt)
                writer.add_scalar('syn_train/acc', np.mean(acc_seq), gcnt)
                loss_seq = []
                acc_seq = []
                sigma_b_loss_seq = []

        scheduler.step()

    if epoch == args.n_epochs - 1:
        epoch = args.n_epochs
        torch.save(get_model_dict(), get_model_path())
        logging.info(f'save train result to {get_model_path()}')
        logging.info('training done!')


def test(args, model):
    test_dataset = lib.data.Codes(args.test_path, train=False)

    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=100,
                                              shuffle=False,
                                              num_workers=8,
                                              worker_init_fn=worker_init_fn)

    assert args.model_path, 'please input model path'

    if args.use_cuda:
        ckpt = torch.load(args.model_path)
    else:
        ckpt = torch.load(args.model_path, map_location=torch.device('cpu'))
    logging.info("load model from {}".format(args.model_path))
    model.load_state_dict(ckpt['model_state_dict'])

    acc_seq = []
    acc_cnt = np.zeros((5, 6))
    acc_tot = np.zeros((5, 6))
    tot = 0
    model.eval()
    # pdb.set_trace()
    SNR = [0, 1, 2, 3, 4]
    for _, (node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b) in tqdm(enumerate(test_loader)):
        if args.use_cuda:
            if args.use_cuda:
                node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b = to_cuda(
                    node_feature, hop_feature, nn_idx_f2v, nn_idx_v2f, efeature_f2v, efeature_v2f, label, sigma_b)

        if len(node_feature.shape) == 3:
            node_feature = node_feature.unsqueeze(-1)
        cur_SNR = node_feature[:, 1, 0, 0]

        with torch.no_grad():
            pred, _ = model(node_feature, hop_feature, nn_idx_f2v,
                            nn_idx_v2f, efeature_f2v, efeature_v2f)

            pred = pred.squeeze().contiguous()
        pred_int = (pred >= 0).long().squeeze()
        label = label.squeeze()

        for csnr in SNR:
            for b in range(6):
                indice = (sigma_b.long() == b) & (abs(cur_SNR - csnr) < 1e-3)
                acc_cnt[csnr][b] += torch.sum(pred_int[indice, :48]
                                              == label[indice, :48]).item()
                acc_tot[csnr][b] += torch.sum(indice) * 48

        # print(i, b)

        # print(
        #     'snr = {} sigma_b = {}  Correct = {}/48'.format(i, b, torch.sum(pred_int[:48] == label[:48]).item()))

        # parameters = list(model.parameters()) + list(emodel_high.parameters())
        # torch.nn.utils.clip_grad_norm(parameters, 1.0)

        all_correct = torch.sum(pred_int[:, :48] == label[:, :48])

        acc_seq.append(all_correct.item())
        tot += np.prod(label.shape) // 2

    print(1 - sum(acc_seq) / tot)
    err_class = 1 - np.divide(acc_cnt, acc_tot)
    print(torch.FloatTensor(err_class))


def residual_link(final_feature, original_feature):
    return final_feature + original_feature[:, :1, :, :]


def main():
    args = parse_args()

    nfeature_dim = 2
    hop_order = 6
    nedge_types = 4
    model = LDPCModel(nfeature_dim, hop_order, nedge_types,
                      aggregator=args.aggregator)

    def get_model_description():
        return str(model)

    logging.info('model {} created'.format(get_model_description()))

    if args.use_cuda:

        model.cuda()

    if args.train:
        subdir = f'train_syn_hop_factor_{args.model_name}_snr_{args.snr}_at_{datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}'
        utils.init_logger('./logs/', subdir, print_log=True)
        logging.info(str(args))
        logdir = f'./tf_logs/{subdir}'
        print(f'logdir = {logdir}')
        writer = SummaryWriter(log_dir=logdir)

        model_dir = f'./model_ldpc/{subdir}'
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        train(args, model, writer, model_dir)

    else:
        test(args, model)


if __name__ == '__main__':
    main()