test

README.md

GOAD

This repository contains a PyTorch implementation of the method presented in "Classification-Based Anomaly Detection for General Data" by Liron Bergman and Yedid Hoshen, ICLR 2020.

Requirements

• Python 3 +
• Pytorch 1.0 +
• Tensorflow 1.8.0 +
• Keras 2.2.0 +
• sklearn 0.19.1 +

Training

To replicate the results of the paper on the tabular-data:

python train_ad_tabular.py --n_rots=64 --n_epoch=25 --d_out=64 --ndf=32 --dataset=kdd 
python train_ad_tabular.py --n_rots=256 --n_epoch=25 --d_out=128 --ndf=128 --dataset=kddrev
python train_ad_tabular.py --n_rots=256 --n_epoch=1 --d_out=32 --ndf=8 --dataset=thyroid
python train_ad_tabular.py --n_rots=256 --n_epoch=1 --d_out=32 --ndf=8 --dataset=arrhythmia 

To replicate the results of the paper on CIFAR10:

python train_ad.py --m=0.1

Citation

If you find this useful, please cite our paper:

@inproceedings{bergman2020goad,
  author    = {Liron Bergman and Yedid Hoshen},
  title     = {Classification-Based Anomaly Detection for General Data},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2020}
}

data_loader.py

import scipy.io
import numpy as np
import pandas as pd
import torchvision.datasets as dset
import torch
import torch.optim as opt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import json
import os
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, Sampler, WeightedRandomSampler




class Data_Loader:

    def __init__(self, n_trains=None):
        self.n_train = n_trains
        self.urls = [
        "http://kdd.ics.uci.edu/databases/kddcup99/kddcup.data_10_percent.gz",
        "http://kdd.ics.uci.edu/databases/kddcup99/kddcup.names"
        ]

    def norm_kdd_data(self, train_real, val_real, val_fake, cont_indices):
        symb_indices = np.delete(np.arange(train_real.shape[1]), cont_indices)
        mus = train_real[:, cont_indices].mean(0)
        sds = train_real[:, cont_indices].std(0)
        sds[sds == 0] = 1

        def get_norm(xs, mu, sd):
            bin_cols = xs[:, symb_indices]
            cont_cols = xs[:, cont_indices]
            cont_cols = np.array([(x - mu) / sd for x in cont_cols])
            return np.concatenate([bin_cols, cont_cols], 1)

        train_real = get_norm(train_real, mus, sds)
        val_real = get_norm(val_real, mus, sds)
        val_fake = get_norm(val_fake, mus, sds)
        return train_real, val_real, val_fake


    def norm_data(self, train_real, val_real, val_fake):
        mus = train_real.mean(0)
        sds = train_real.std(0)
        sds[sds == 0] = 1

        def get_norm(xs, mu, sd):
            return np.array([(x - mu) / sd for x in xs])

        train_real = get_norm(train_real, mus, sds)
        val_real = get_norm(val_real, mus, sds)
        val_fake = get_norm(val_fake, mus, sds)
        return train_real, val_real, val_fake

    def norm(self, data, mu=1):
        return 2 * (data / 255.) - mu

    def get_dataset(self, dataset_name, c_percent=None, true_label=1, input_length=None):
        if dataset_name == 'cifar10':
            return self.load_data_CIFAR10(true_label)
        if dataset_name == 'kdd':
            return self.KDD99_train_valid_data()
        if dataset_name == 'kddrev':
            return self.KDD99Rev_train_valid_data()
        if dataset_name == 'thyroid':
            return self.Thyroid_train_valid_data()
        if dataset_name == 'arrhythmia':
            return self.Arrhythmia_train_valid_data()
        if dataset_name == 'ckdd':
            return self.contaminatedKDD99_train_valid_data(c_percent)
        else:
            return self.NAB_data(dataset_name, input_length)


    def load_data_CIFAR10(self, true_label):
        root = './data'
        if not os.path.exists(root):
            os.mkdir(root)

        trainset = dset.CIFAR10(root, train=True, download=True)
        train_data = np.array(trainset.data)
        train_labels = np.array(trainset.targets)

        testset = dset.CIFAR10(root, train=False, download=True)
        test_data = np.array(testset.data)
        test_labels = np.array(testset.targets)

        train_data = train_data[np.where(train_labels == true_label)]
        x_train = self.norm(np.asarray(train_data, dtype='float32'))
        x_test = self.norm(np.asarray(test_data, dtype='float32'))
        return x_train, x_test, test_labels

    def NAB_data(self, file_name, input_length):
        ds = MySeriesDataset("NAB", data_file=file_name, size=20)
        samples = ds.chunks
        labels = ds.labels

        norm_samples = samples[labels == 0]  # 3679 norm
        anom_samples = samples[labels == 1]  # 93 anom

        n_train = len(norm_samples) // 2
        x_train = norm_samples[:n_train]  # 1839 train

        val_real = norm_samples[n_train:]
        val_fake = anom_samples


        return x_train.numpy(), val_real.numpy(), val_fake.numpy()




    def Thyroid_train_valid_data(self):
        data = scipy.io.loadmat("data/thyroid.mat")
        samples = data['X']  # 3772
        labels = ((data['y']).astype(np.int32)).reshape(-1)

        norm_samples = samples[labels == 0]  # 3679 norm
        anom_samples = samples[labels == 1]  # 93 anom

        n_train = len(norm_samples) // 2
        x_train = norm_samples[:n_train]  # 1839 train

        val_real = norm_samples[n_train:]
        val_fake = anom_samples
        return self.norm_data(x_train, val_real, val_fake)


    def Arrhythmia_train_valid_data(self):
        data = scipy.io.loadmat("data/arrhythmia.mat")
        samples = data['X']  # 518
        labels = ((data['y']).astype(np.int32)).reshape(-1)

        norm_samples = samples[labels == 0]  # 452 norm
        anom_samples = samples[labels == 1]  # 66 anom

        n_train = len(norm_samples) // 2
        x_train = norm_samples[:n_train]  # 226 train

        val_real = norm_samples[n_train:]
        val_fake = anom_samples
        return self.norm_data(x_train, val_real, val_fake)


    def KDD99_preprocessing(self):
        df_colnames = pd.read_csv(self.urls[1], skiprows=1, sep=':', names=['f_names', 'f_types'])
        df_colnames.loc[df_colnames.shape[0]] = ['status', ' symbolic.']
        df = pd.read_csv(self.urls[0], header=None, names=df_colnames['f_names'].values)
        df_symbolic = df_colnames[df_colnames['f_types'].str.contains('symbolic.')]
        df_continuous = df_colnames[df_colnames['f_types'].str.contains('continuous.')]
        samples = pd.get_dummies(df.iloc[:, :-1], columns=df_symbolic['f_names'][:-1])

        smp_keys = samples.keys()
        cont_indices = []
        for cont in df_continuous['f_names']:
            cont_indices.append(smp_keys.get_loc(cont))

        labels = np.where(df['status'] == 'normal.', 1, 0)
        return np.array(samples), np.array(labels), cont_indices


    def KDD99_train_valid_data(self):
        samples, labels, cont_indices = self.KDD99_preprocessing()
        anom_samples = samples[labels == 1]  # norm: 97278

        norm_samples = samples[labels == 0]  # attack: 396743

        n_norm = norm_samples.shape[0]
        ranidx = np.random.permutation(n_norm)
        n_train = n_norm // 2

        x_train = norm_samples[ranidx[:n_train]]
        norm_test = norm_samples[ranidx[n_train:]]

        val_real = norm_test
        val_fake = anom_samples
        return self.norm_kdd_data(x_train, val_real, val_fake, cont_indices)


    def KDD99Rev_train_valid_data(self):
        samples, labels, cont_indices = self.KDD99_preprocessing()

        norm_samples = samples[labels == 1]  # norm: 97278

        # Randomly draw samples labeled as 'attack'
        # so that the ratio btw norm:attack will be 4:1
        # len(anom) = 24,319
        anom_samples = samples[labels == 0]  # attack: 396743

        rp = np.random.permutation(len(anom_samples))
        rp_cut = rp[:24319]
        anom_samples = anom_samples[rp_cut]  # attack:24319

        n_norm = norm_samples.shape[0]
        ranidx = np.random.permutation(n_norm)
        n_train = n_norm // 2

        x_train = norm_samples[ranidx[:n_train]]
        norm_test = norm_samples[ranidx[n_train:]]

        val_real = norm_test
        val_fake = anom_samples
        return self.norm_kdd_data(x_train, val_real, val_fake, cont_indices)


    def contaminatedKDD99_train_valid_data(self, c_percent):
        samples, labels, cont_indices = self.KDD99_preprocessing()

        ranidx = np.random.permutation(len(samples))
        n_test = len(samples)//2
        x_test = samples[ranidx[:n_test]]
        y_test = labels[ranidx[:n_test]]

        x_train = samples[ranidx[n_test:]]
        y_train = labels[ranidx[n_test:]]

        norm_samples = x_train[y_train == 0]  # attack: 396743
        anom_samples = x_train[y_train == 1]  # norm: 97278
        n_contaminated = int((c_percent/100)*len(anom_samples))

        rpc = np.random.permutation(n_contaminated)
        x_train = np.concatenate([norm_samples, anom_samples[rpc]])

        val_real = x_test[y_test == 0]
        val_fake = x_test[y_test == 1]
        return self.norm_kdd_data(x_train, val_real, val_fake, cont_indices)



class MySeriesDataset:
    def __init__(self, dname, data_file, size, step=1):
        label_file = "combined_labels"
        with open(os.path.join("data", dname, "labels", label_file + '.json'), 'r') as f:
            self.label_json = json.load(f)

        data_file_path = self.getRelativePath(data_file)
        assert data_file_path is not -1
        data_df = pd.read_csv(os.path.join("data", dname, "data", data_file_path))
        data_df['timestamp'] = pd.to_datetime(data_df['timestamp'])
        data_df['stand_value'] = standardization(data_df['value'])
        anomalies_ts = [a_ts.replace(".000000", "") for a_ts in self.label_json[data_file_path]]
        anomalies = data_df[data_df['timestamp'].isin(anomalies_ts)]
        # anomalies['timestamp'] = pd.to_datetime(anomalies['timestamp'])

        self.chunks = torch.FloatTensor(data_df['value']).unfold(0, size, step)
        # self.chunks = self.chunks.view(-1, size)
        self.df = data_df
        self.anomalies = anomalies

        '''
        dataframe没有unfold函数，但是有rolling滑窗函数，rolling如果要实现自定义函数需要传数值类型的列，
        但是真正的数据用不到，而是通过传入Series的index判断某条数据是不是异常数据，最后生成的chunks是size+1, 
        所以用size为period做rolling运算会多一个label数据。
        '''
        self.labels = torch.tensor(data_df['value'].rolling(size).apply(
            lambda x:self.ifAnomalySeries(x, anomalies.index.tolist())).dropna().tolist(), dtype=torch.int64)#[:-1]


    def ifAnomalySeries(self, windowed_data, anomalies_idx):
        if True in [x in anomalies_idx for x in windowed_data.index.tolist()]:
            return 1
        else:
            return 0


    def getRelativePath(self, data_file):
        # 输入目标csv文件名字， 在label的json中找到了返回这个csv文件的完整路径，找不到返回-1
        for file in list(self.label_json.keys()):
            if file.rfind(data_file) is not -1:
                return file
        return -1


    def __len__(self):
        return self.chunks.size(0)


    def __getitem__(self, i):
        x = self.chunks[i, :, :] # 倒数第一个是要预测的label，因此从第一个取到倒数第二个
        # y = self.chunks[i, :, -1:].squeeze(-1) # 取倒数第一个作为label
        anomaly_label = self.labels[i]
        return x, anomaly_label


def normalization(data):
    return (data - data.min())/(data.max() - data.min())


def standardization(data):
    return (data - data.mean())/(data.std())

fcnet.py

import torch.nn as nn
import torch.nn.init as init
import numpy as np

def weights_init(m):
    classname = m.__class__.__name__
    if isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight, gain=np.sqrt(2.0))
    elif classname.find('Conv') != -1:
        init.xavier_normal_(m.weight, gain=np.sqrt(2.0))
    elif classname.find('Linear') != -1:
        init.eye_(m.weight)
    elif classname.find('Emb') != -1:
        init.normal(m.weight, mean=0, std=0.01)

class netC5(nn.Module):
    def __init__(self, d, ndf, nc):
        super(netC5, self).__init__()
        self.trunk = nn.Sequential(
        nn.Conv1d(d, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        )
        self.head = nn.Sequential(
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, nc, kernel_size=1, bias=True),
        )


    def forward(self, input):
        tc = self.trunk(input)
        ce = self.head(tc)
        return tc, ce


class netC1(nn.Module):
    def __init__(self, d, ndf, nc): # d:d_out, ndf:ndf, nc:n_rots
        super(netC1, self).__init__()
        self.trunk = nn.Sequential(
        nn.Conv1d(d, ndf, kernel_size=1, bias=False),
        )
        self.head = nn.Sequential(
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, nc, kernel_size=1, bias=True),

        )
        # self.pool = nn.AdaptiveAvgPool1d(1)
        # self.liner = nn.Linear(ndf, nc, bias=True)

    def forward(self, input):
        tc = self.trunk(input)

        ce = self.head(tc)
        # ce = self.pool(ce)
        # ce = ce.view(ce.shape[0], -1)
        # ce = self.liner(ce)
        return tc, ce

fcnet_test.py

import torch.nn as nn
import torch.nn.init as init
import numpy as np

def weights_init(m):
    classname = m.__class__.__name__
    if isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight, gain=np.sqrt(2.0))
    elif classname.find('Conv') != -1:
        init.xavier_normal_(m.weight, gain=np.sqrt(2.0))
    elif classname.find('Linear') != -1:
        init.eye_(m.weight)
    elif classname.find('Emb') != -1:
        init.normal(m.weight, mean=0, std=0.01)

class netC5(nn.Module):
    def __init__(self, d, ndf, nc):
        super(netC5, self).__init__()
        self.trunk = nn.Sequential(
        nn.Conv1d(d, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, ndf, kernel_size=1, bias=False),
        )
        self.head = nn.Sequential(
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, nc, kernel_size=1, bias=True),
        )


    def forward(self, input):
        tc = self.trunk(input)
        ce = self.head(tc)
        return tc, ce


class netC1(nn.Module):
    def __init__(self, d, ndf, nc): # d:d_out, ndf:ndf, nc:n_rots
        super(netC1, self).__init__()
        self.trunk = nn.Sequential(
        nn.Conv1d(d, ndf, kernel_size=1, bias=False),
        )
        self.head = nn.Sequential(
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv1d(ndf, nc, kernel_size=1, bias=True),

        )
        # self.pool = nn.AdaptiveAvgPool1d(1)
        # self.liner = nn.Linear(ndf, nc, bias=True)

    def forward(self, input):
        tc = self.trunk(input)

        ce = self.head(tc)
        # ce = self.pool(ce)
        # ce = ce.view(ce.shape[0], -1)
        # ce = self.liner(ce)
        return tc, ce

opt_tc.py

import torch.utils.data
import numpy as np
import torch
import torch.utils.data
from torch.backends import cudnn
from wideresnet import WideResNet
from sklearn.metrics import roc_auc_score

cudnn.benchmark = True

def tc_loss(zs, m):
    means = zs.mean(0).unsqueeze(0)
    res = ((zs.unsqueeze(2) - means.unsqueeze(1)) ** 2).sum(-1)
    pos = torch.diagonal(res, dim1=1, dim2=2)
    offset = torch.diagflat(torch.ones(zs.size(1))).unsqueeze(0).cuda() * 1e6
    neg = (res + offset).min(-1)[0]
    loss = torch.clamp(pos + m - neg, min=0).mean()
    return loss


class TransClassifier():
    def __init__(self, num_trans, args):
        self.n_trans = num_trans
        self.args = args
        self.netWRN = WideResNet(self.args.depth, num_trans, self.args.widen_factor).cuda()
        self.optimizer = torch.optim.Adam(self.netWRN.parameters())


    def fit_trans_classifier(self, x_train, x_test, y_test):
        print("Training")
        self.netWRN.train()
        bs = self.args.batch_size
        N, sh, sw, nc = x_train.shape
        n_rots = self.n_trans
        m = self.args.m
        celoss = torch.nn.CrossEntropyLoss()
        ndf = 256

        for epoch in range(self.args.epochs):
            rp = np.random.permutation(N//n_rots)
            rp = np.concatenate([np.arange(n_rots) + rp[i]*n_rots for i in range(len(rp))])
            assert len(rp) == N
            all_zs = torch.zeros((len(x_train), ndf)).cuda()
            diffs_all = []

            for i in range(0, len(x_train), bs):
                batch_range = min(bs, len(x_train) - i)
                idx = np.arange(batch_range) + i
                xs = torch.from_numpy(x_train[rp[idx]]).float().cuda()
                zs_tc, zs_ce = self.netWRN(xs)

                all_zs[idx] = zs_tc
                train_labels = torch.from_numpy(np.tile(np.arange(n_rots), batch_range//n_rots)).long().cuda()
                zs = torch.reshape(zs_tc, (batch_range//n_rots, n_rots, ndf))

                means = zs.mean(0).unsqueeze(0)
                diffs = -((zs.unsqueeze(2).detach().cpu().numpy() - means.unsqueeze(1).detach().cpu().numpy()) ** 2).sum(-1)
                diffs_all.append(torch.diagonal(torch.tensor(diffs), dim1=1, dim2=2))

                tc = tc_loss(zs, m)
                ce = celoss(zs_ce, train_labels)
                if self.args.reg:
                    loss = ce + self.args.lmbda * tc + 10 *(zs*zs).mean()
                else:
                    loss = ce + self.args.lmbda * tc
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            self.netWRN.eval()
            all_zs = torch.reshape(all_zs, (N//n_rots, n_rots, ndf))
            means = all_zs.mean(0, keepdim=True)


            with torch.no_grad():
                batch_size = bs
                val_probs_rots = np.zeros((len(y_test), self.n_trans))
                for i in range(0, len(x_test), batch_size):
                    batch_range = min(batch_size, len(x_test) - i)
                    idx = np.arange(batch_range) + i
                    xs = torch.from_numpy(x_test[idx]).float().cuda()

                    zs, fs = self.netWRN(xs)
                    zs = torch.reshape(zs, (batch_range // n_rots, n_rots, ndf))

                    diffs = ((zs.unsqueeze(2) - means) ** 2).sum(-1)
                    diffs_eps = self.args.eps * torch.ones_like(diffs)
                    diffs = torch.max(diffs, diffs_eps)
                    logp_sz = torch.nn.functional.log_softmax(-diffs, dim=2)

                    zs_reidx = np.arange(batch_range // n_rots) + i // n_rots
                    val_probs_rots[zs_reidx] = -torch.diagonal(logp_sz, 0, 1, 2).cpu().data.numpy()

                val_probs_rots = val_probs_rots.sum(1)
                print("Epoch:", epoch, ", AUC: ", roc_auc_score(y_test, -val_probs_rots))

opt_tc_tabular.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import fcnet as model
from sklearn.metrics import precision_recall_fscore_support as prf

def tc_loss(zs, m):
    # zs: (batch_size, class_num, ndf) -> (batch_size, ndf, class_num)
    means = zs.mean(0).unsqueeze(0)
    res = ((zs.unsqueeze(2) - means.unsqueeze(1)) ** 2).sum(-1)
    pos = torch.diagonal(res, dim1=1, dim2=2)
    offset = torch.diagflat(torch.ones(zs.size(1))).unsqueeze(0).cuda() * 1e6
    neg = (res + offset).min(-1)[0]
    loss = torch.clamp(pos + m - neg, min=0).mean()
    return loss

def f_score(scores, labels, ratio):
    thresh = np.percentile(scores, ratio)
    y_pred = (scores >= thresh).astype(int)
    y_true = labels.astype(int)
    precision, recall, f_score, support = prf(y_true, y_pred, average='binary')
    return f_score


class TransClassifierTabular():
    def __init__(self, args):
        self.ds = args.dataset
        self.m = args.m
        self.lmbda = args.lmbda
        self.batch_size = args.batch_size
        self.ndf = args.ndf
        self.n_rots = args.n_rots
        self.d_out = args.d_out
        self.eps = args.eps

        self.n_epoch = args.n_epoch
        if False:#args.dataset == "thyroid" or args.dataset == "arrhythmia":
            self.netC = model.netC1(self.d_out, self.ndf, self.n_rots).cuda()
        else:
            self.netC = model.netC5(self.d_out, self.ndf, self.n_rots).cuda()
        model.weights_init(self.netC)
        self.optimizerC = optim.Adam(self.netC.parameters(), lr=args.lr, betas=(0.5, 0.999))


    def fit_trans_classifier(self, train_xs, x_test, y_test, ratio):
        labels = torch.arange(self.n_rots).unsqueeze(0).expand((self.batch_size, self.n_rots)).long().cuda()
        celoss = nn.CrossEntropyLoss()
        print('Training')
        for epoch in range(self.n_epoch):
            self.netC.train()
            rp = np.random.permutation(len(train_xs))
            n_batch = 0
            sum_zs = torch.zeros((self.ndf, self.n_rots)).cuda()

            for i in range(0, len(train_xs), self.batch_size):
                self.netC.zero_grad()
                batch_range = min(self.batch_size, len(train_xs) - i) # i就是已经迭代完的数据的个数
                train_labels = labels
                if batch_range == len(train_xs) - i:
                    train_labels = torch.arange(self.n_rots).unsqueeze(0).expand((len(train_xs) - i, self.n_rots)).long().cuda()
                idx = np.arange(batch_range) + i # 依次取数据
                xs = torch.from_numpy(train_xs[rp[idx]]).float().cuda()
                tc_zs, ce_zs = self.netC(xs)
                sum_zs = sum_zs + tc_zs.mean(0)
                tc_zs = tc_zs.permute(0, 2, 1) #(batch_size, class_num, ndf) -> (batch_size, ndf, class_num)

                loss_ce = celoss(ce_zs, train_labels)
                print(loss_ce.item())
                er = loss_ce + self.lmbda * tc_loss(tc_zs, self.m)
                er.backward()
                self.optimizerC.step()
                n_batch += 1

            means = sum_zs.t() / n_batch
            means = means.unsqueeze(0)
            self.netC.eval()

            with torch.no_grad():
                val_probs_rots = np.zeros((len(y_test), self.n_rots))
                for i in range(0, len(x_test), self.batch_size):
                    batch_range = min(self.batch_size, len(x_test) - i)
                    idx = np.arange(batch_range) + i
                    xs = torch.from_numpy(x_test[idx]).float().cuda()
                    zs, fs = self.netC(xs)
                    zs = zs.permute(0, 2, 1)
                    diffs = ((zs.unsqueeze(2) - means) ** 2).sum(-1)

                    diffs_eps = self.eps * torch.ones_like(diffs)
                    diffs = torch.max(diffs, diffs_eps)
                    logp_sz = torch.nn.functional.log_softmax(-diffs, dim=2)

                    val_probs_rots[idx] = -torch.diagonal(logp_sz, 0, 1, 2).cpu().data.numpy()

                val_probs_rots = val_probs_rots.sum(1)
                f1_score = f_score(val_probs_rots, y_test, ratio)
                print("Epoch:", epoch, ", fscore: ", f1_score)
        return f1_score

opt_tc_tabular_backup.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import fcnet as model
from sklearn.metrics import precision_recall_fscore_support as prf

def tc_loss(zs, m):
    # zs: (batch_size, class_num, ndf) -> (batch_size, ndf, class_num)
    means = zs.mean(0).unsqueeze(0)
    res = ((zs.unsqueeze(2) - means.unsqueeze(1)) ** 2).sum(-1)
    pos = torch.diagonal(res, dim1=1, dim2=2)
    offset = torch.diagflat(torch.ones(zs.size(1))).unsqueeze(0).cuda() * 1e6
    neg = (res + offset).min(-1)[0]
    loss = torch.clamp(pos + m - neg, min=0).mean()
    return loss

def f_score(scores, labels, ratio):
    thresh = np.percentile(scores, ratio)
    y_pred = (scores >= thresh).astype(int)
    y_true = labels.astype(int)
    precision, recall, f_score, support = prf(y_true, y_pred, average='binary')
    return f_score


class TransClassifierTabular():
    def __init__(self, args):
        self.ds = args.dataset
        self.m = args.m
        self.lmbda = args.lmbda
        self.batch_size = args.batch_size
        self.ndf = args.ndf
        self.n_rots = args.n_rots
        self.d_out = args.d_out
        self.eps = args.eps

        self.n_epoch = args.n_epoch
        if args.dataset == "thyroid" or args.dataset == "arrhythmia":
            self.netC = model.netC1(self.d_out, self.ndf, self.n_rots).cuda()
        else:
            self.netC = model.netC5(self.d_out, self.ndf, self.n_rots).cuda()
        model.weights_init(self.netC)
        self.optimizerC = optim.Adam(self.netC.parameters(), lr=args.lr, betas=(0.5, 0.999))


    def fit_trans_classifier(self, train_xs, x_test, y_test, ratio):
        labels = torch.arange(self.n_rots).unsqueeze(0).expand((self.batch_size, self.n_rots)).long().cuda()
        celoss = nn.CrossEntropyLoss()
        print('Training')
        for epoch in range(self.n_epoch):
            self.netC.train()
            rp = np.random.permutation(len(train_xs))
            n_batch = 0
            sum_zs = torch.zeros((self.ndf, self.n_rots)).cuda()

            for i in range(0, len(train_xs), self.batch_size):
                self.netC.zero_grad()
                batch_range = min(self.batch_size, len(train_xs) - i) # i就是已经迭代完的数据的个数
                train_labels = labels
                if batch_range == len(train_xs) - i:
                    train_labels = torch.arange(self.n_rots).unsqueeze(0).expand((len(train_xs) - i, self.n_rots)).long().cuda()
                idx = np.arange(batch_range) + i # 依次取数据
                xs = torch.from_numpy(train_xs[rp[idx]]).float().cuda()
                tc_zs, ce_zs = self.netC(xs)
                sum_zs = sum_zs + tc_zs.mean(0)
                tc_zs = tc_zs.permute(0, 2, 1) #(batch_size, class_num, ndf) -> (batch_size, ndf, class_num)

                loss_ce = celoss(ce_zs, train_labels)
                er = loss_ce + self.lmbda * tc_loss(tc_zs, self.m)
                er.backward()
                self.optimizerC.step()
                n_batch += 1

            means = sum_zs.t() / n_batch
            means = means.unsqueeze(0)
            self.netC.eval()

            with torch.no_grad():
                val_probs_rots = np.zeros((len(y_test), self.n_rots))
                for i in range(0, len(x_test), self.batch_size):
                    batch_range = min(self.batch_size, len(x_test) - i)
                    idx = np.arange(batch_range) + i
                    xs = torch.from_numpy(x_test[idx]).float().cuda()
                    zs, fs = self.netC(xs)
                    zs = zs.permute(0, 2, 1)
                    diffs = ((zs.unsqueeze(2) - means) ** 2).sum(-1)

                    diffs_eps = self.eps * torch.ones_like(diffs)
                    diffs = torch.max(diffs, diffs_eps)
                    logp_sz = torch.nn.functional.log_softmax(-diffs, dim=2)

                    val_probs_rots[idx] = -torch.diagonal(logp_sz, 0, 1, 2).cpu().data.numpy()

                val_probs_rots = val_probs_rots.sum(1)
                f1_score = f_score(val_probs_rots, y_test, ratio)
                print("Epoch:", epoch, ", fscore: ", f1_score)
        return f1_score

opt_tc_tabular_test.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import fcnet as model
from sklearn.metrics import precision_recall_fscore_support as prf

def tc_loss(zs, m):
    means = zs.mean(0).unsqueeze(0)
    res = ((zs.unsqueeze(2) - means.unsqueeze(1)) ** 2).sum(-1)
    pos = torch.diagonal(res, dim1=1, dim2=2)
    offset = torch.diagflat(torch.ones(zs.size(1))).unsqueeze(0).cuda() * 1e6
    neg = (res + offset).min(-1)[0]
    loss = torch.clamp(pos + m - neg, min=0).mean()
    return loss

def f_score(scores, labels, ratio):
    thresh = np.percentile(scores, ratio)
    y_pred = (scores >= thresh).astype(int)
    y_true = labels.astype(int)
    precision, recall, f_score, support = prf(y_true, y_pred, average='binary')
    return f_score


class TransClassifierTabular():
    def __init__(self, args):
        self.ds = args.dataset
        self.m = args.m
        self.lmbda = args.lmbda
        self.batch_size = args.batch_size
        self.ndf = args.ndf
        self.n_rots = args.n_rots
        self.d_out = args.d_out
        self.eps = args.eps

        self.n_epoch = args.n_epoch
        if args.dataset == "thyroid" or args.dataset == "arrhythmia":
            self.netC = model.netC1(self.d_out, self.ndf, self.n_rots).cuda()
        else:
            self.netC = model.netC5(self.d_out, self.ndf, self.n_rots).cuda()
        model.weights_init(self.netC)
        self.optimizerC = optim.Adam(self.netC.parameters(), lr=args.lr, betas=(0.5, 0.999))


    def fit_trans_classifier(self, train_xs, x_test, y_test, ratio):
        labels = torch.cat([torch.arange(self.n_rots) for i in range(int(train_xs.shape[0]/self.n_rots))]).long().cuda()
        celoss = nn.CrossEntropyLoss()
        print('Training')
        for epoch in range(self.n_epoch):
            self.netC.train()
            rp = np.random.permutation(len(train_xs))
            n_batch = 0
            sum_zs = torch.zeros((self.ndf, self.n_rots)).cuda()

            for i in range(0, len(train_xs), self.batch_size):
                self.netC.zero_grad()
                batch_range = min(self.batch_size, len(train_xs) - i*self.batch_size)
                train_labels = labels
                if batch_range != self.batch_size:
                    train_labels = torch.arange(batch_range).unsqueeze(0).long().cuda()
                idx = np.arange(batch_range*self.n_rots) + i
                xs = train_xs[rp[idx]].float().cuda()
                tc_zs, ce_zs = self.netC(xs)
                sum_zs = sum_zs + tc_zs.mean(0)
                tc_zs = tc_zs.permute(0, 2, 1) #(batch_size, class_num, ndf) -> (batch_size, ndf, class_num)

                loss_ce = celoss(ce_zs, train_labels)
                er = self.lmbda * tc_loss(tc_zs, self.m) + loss_ce
                er.backward()
                self.optimizerC.step()
                n_batch += 1

            means = sum_zs.t() / n_batch
            means = means.unsqueeze(0)
            self.netC.eval()

            with torch.no_grad():
                val_probs_rots = np.zeros((len(y_test), self.n_rots))
                for i in range(0, len(x_test), self.batch_size):
                    batch_range = min(self.batch_size, len(x_test) - i)
                    idx = np.arange(batch_range) + i
                    xs = torch.from_numpy(x_test[idx]).float().cuda()
                    zs, fs = self.netC(xs)
                    zs = zs.permute(0, 2, 1)
                    diffs = ((zs.unsqueeze(2) - means) ** 2).sum(-1)

                    diffs_eps = self.eps * torch.ones_like(diffs)
                    diffs = torch.max(diffs, diffs_eps)
                    logp_sz = torch.nn.functional.log_softmax(-diffs, dim=2)

                    val_probs_rots[idx] = -torch.diagonal(logp_sz, 0, 1, 2).cpu().data.numpy()

                val_probs_rots = val_probs_rots.sum(1)
                f1_score = f_score(val_probs_rots, y_test, ratio)
                print("Epoch:", epoch, ", fscore: ", f1_score)
        return f1_score

requirements.txt

torch==1.4.0
tensorflow>=1.15.2
scikit-learn==0.19.1

train_ad.py

import argparse
import transformations as ts
import opt_tc as tc
import numpy as np
from data_loader import Data_Loader

def transform_data(data, trans):
    trans_inds = np.tile(np.arange(trans.n_transforms), len(data))
    trans_data = trans.transform_batch(np.repeat(np.array(data), trans.n_transforms, axis=0), trans_inds)
    return trans_data, trans_inds

def load_trans_data(args, trans):
    dl = Data_Loader()
    x_train, x_test, y_test = dl.get_dataset(args.dataset, true_label=args.class_ind)
    x_train_trans, labels = transform_data(x_train, trans)
    x_test_trans, _ = transform_data(x_test, trans)
    x_test_trans, x_train_trans = x_test_trans.transpose(0, 3, 1, 2), x_train_trans.transpose(0, 3, 1, 2)
    y_test = np.array(y_test) == args.class_ind
    return x_train_trans, x_test_trans, y_test


def train_anomaly_detector(args):
    transformer = ts.get_transformer(args.type_trans)
    x_train, x_test, y_test = load_trans_data(args, transformer)
    tc_obj = tc.TransClassifier(transformer.n_transforms, args)
    tc_obj.fit_trans_classifier(x_train, x_test, y_test)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Wide Residual Networks')
    # Model options
    parser.add_argument('--depth', default=10, type=int)
    parser.add_argument('--widen-factor', default=4, type=int)

    # Training options
    parser.add_argument('--batch_size', default=288, type=int)
    parser.add_argument('--lr', '--learning-rate', default=0.001, type=float)
    parser.add_argument('--epochs', default=16, type=int)

    # Trans options
    parser.add_argument('--type_trans', default='complicated', type=str)

    # CT options
    parser.add_argument('--lmbda', default=0.1, type=float)
    parser.add_argument('--m', default=1, type=float)
    parser.add_argument('--reg', default=True, type=bool)
    parser.add_argument('--eps', default=0, type=float)

    # Exp options
    parser.add_argument('--class_ind', default=1, type=int)
    parser.add_argument('--dataset', default='cifar10', type=str)
    args = parser.parse_args()

    for i in range(10):
        args.class_ind = i
        print("Dataset: CIFAR10")
        print("True Class:", args.class_ind)
        train_anomaly_detector(args)

train_ad_NAB.py

import numpy as np
from data_loader import Data_Loader
import opt_tc_tabular as tc
import argparse
import torch
import torch.optim as opt
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import json
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, Sampler, WeightedRandomSampler

def load_trans_data(args):
    dl = Data_Loader()
    # val_fake val中的anomalous data, val_real val中的normal data
    train_real, val_real, val_fake = dl.get_dataset(args.dataset, args.c_pr)
    y_test_fscore = np.concatenate([np.zeros(len(val_real)), np.ones(len(val_fake))])
    ratio = 100.0 * len(val_real) / (len(val_real) + len(val_fake)) # 数据集中正常数据的比率

    n_train, n_dims = train_real.shape
    rots = np.random.randn(args.n_rots, n_dims, args.d_out) #(tranformations_num, attribute_num, 4)

    print('Calculating transforms')
    #每一种transformation就是一个weight，与input做矩阵乘法运算。
    x_train = np.stack([train_real.dot(rot) for rot in rots], 2)    # train_real:(data_num, attribute_num); rot:(attribute_num, d_out)
    val_real_xs = np.stack([val_real.dot(rot) for rot in rots], 2)
    val_fake_xs = np.stack([val_fake.dot(rot) for rot in rots], 2)
    x_test = np.concatenate([val_real_xs, val_fake_xs])
    return x_train, x_test, y_test_fscore, ratio


def train_anomaly_detector(args):
    x_train, x_test, y_test, ratio = load_trans_data(args)
    tc_obj = tc.TransClassifierTabular(args)
    f_score = tc_obj.fit_trans_classifier(x_train, x_test, y_test, ratio)
    return f_score

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--lr', default=0.001, type=float)
    parser.add_argument('--n_rots', default=80, type=int)
    parser.add_argument('--batch_size', default=64, type=int)
    parser.add_argument('--n_epoch', default=1, type=int)
    parser.add_argument('--d_out', default=18, type=int)
    parser.add_argument('--dataset', default='thyroid', type=str)   # arrhythmia  thyroid
    parser.add_argument('--exp', default='affine', type=str)
    parser.add_argument('--c_pr', default=0, type=int)
    parser.add_argument('--true_label', default=1, type=int)
    parser.add_argument('--ndf', default=8, type=int)
    parser.add_argument('--m', default=1, type=float)
    parser.add_argument('--lmbda', default=0.1, type=float)
    parser.add_argument('--eps', default=0, type=float)
    parser.add_argument('--n_iters', default=500, type=int)

    args = parser.parse_args()
    print("Dataset: ", args.dataset)

    if args.dataset == 'thyroid' or args.dataset == 'arrhythmia':
        n_iters = args.n_iters
        f_scores = np.zeros(n_iters)
        for i in range(n_iters):
            f_scores[i] = train_anomaly_detector(args)
        print("AVG f1_score", f_scores.mean())
    else:
        train_anomaly_detector(args)

train_ad_tabular.py

import numpy as np
from data_loader import Data_Loader
import opt_tc_tabular as tc
import argparse

def load_trans_data(args):
    dl = Data_Loader()
    # val_fake val中的anomalous data, val_real val中的normal data
    train_real, val_real, val_fake = dl.get_dataset(args.dataset, args.c_pr)
    y_test_fscore = np.concatenate([np.zeros(len(val_real)), np.ones(len(val_fake))])
    ratio = 100.0 * len(val_real) / (len(val_real) + len(val_fake)) # 数据集中正常数据的比率

    n_train, n_dims = train_real.shape
    rots = np.random.randn(args.n_rots, n_dims, args.d_out) #(tranformations_num, attribute_num, 4)

    print('Calculating transforms')
    #每一种transformation就是一个weight，与input做矩阵乘法运算。
    x_train = np.stack([train_real.dot(rot) for rot in rots], 2)    # train_real:(data_num, attribute_num); rot:(attribute_num, d_out)
    val_real_xs = np.stack([val_real.dot(rot) for rot in rots], 2)
    val_fake_xs = np.stack([val_fake.dot(rot) for rot in rots], 2)
    x_test = np.concatenate([val_real_xs, val_fake_xs])
    return x_train, x_test, y_test_fscore, ratio


def train_anomaly_detector(args):
    x_train, x_test, y_test, ratio = load_trans_data(args)
    tc_obj = tc.TransClassifierTabular(args)
    f_score = tc_obj.fit_trans_classifier(x_train, x_test, y_test, ratio)
    return f_score

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--lr', default=0.05, type=float)
    parser.add_argument('--n_rots', default=128, type=int)
    parser.add_argument('--batch_size', default=64, type=int)
    parser.add_argument('--n_epoch', default=10, type=int)# Twitter_volume_CVS
    parser.add_argument('--d_out', default=32, type=int) # rds_cpu_utilization_cc0c53
    parser.add_argument('--dataset', default='exchange-3_cpc_results', type=str)   # arrhythmia  thyroid
    parser.add_argument('--exp', default='affine', type=str)
    parser.add_argument('--c_pr', default=0, type=int)
    parser.add_argument('--true_label', default=1, type=int)
    parser.add_argument('--ndf', default=8, type=int)
    parser.add_argument('--m', default=1, type=float)
    parser.add_argument('--lmbda', default=0.1, type=float)
    parser.add_argument('--eps', default=0, type=float)
    parser.add_argument('--n_iters', default=500, type=int)

    args = parser.parse_args()
    print("Dataset: ", args.dataset)

    if False:#args.dataset == 'thyroid' or args.dataset == 'arrhythmia':
        n_iters = args.n_iters
        f_scores = np.zeros(n_iters)
        for i in range(n_iters):
            f_scores[i] = train_anomaly_detector(args)
        print("AVG f1_score", f_scores.mean())
    else:
        train_anomaly_detector(args)

train_ad_tabular_test.py

import numpy as np
from data_loader import Data_Loader
import opt_tc_tabular_test as tc
import argparse
import torch

def load_trans_data(args):
    dl = Data_Loader()
    # val_fake val中的anomalous data, val_real val中的normal data
    train_real, val_real, val_fake = dl.get_dataset(args.dataset, args.c_pr)
    y_test_fscore = np.concatenate([np.zeros(len(val_real)), np.ones(len(val_fake))])
    ratio = 100.0 * len(val_real) / (len(val_real) + len(val_fake))

    n_train, n_dims = train_real.shape
    rots = np.random.randn(args.n_rots, n_dims, args.d_out) #(tranformations_num, attribute_num, 4)

    print('Calculating transforms')
    #每一种transformation就是一个weight，与input做矩阵乘法运算。
    x_train = torch.cat([torch.tensor(train_real.dot(rot)) for rot in rots])    # train_real:(data_num, attribute_num); rot:(attribute_num, d_out)

    val_real_xs = torch.cat([torch.tensor(val_real.dot(rot)) for rot in rots])
    val_fake_xs = torch.cat([torch.tensor(val_fake.dot(rot)) for rot in rots])
    x_test = torch.cat([val_real_xs, val_fake_xs])
    return x_train, x_test, y_test_fscore, ratio


def train_anomaly_detector(args):
    x_train, x_test, y_test, ratio = load_trans_data(args)
    tc_obj = tc.TransClassifierTabular(args)
    f_score = tc_obj.fit_trans_classifier(x_train, x_test, y_test, ratio)
    return f_score

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--lr', default=0.001, type=float)
    parser.add_argument('--n_rots', default=32, type=int)
    parser.add_argument('--batch_size', default=64, type=int)
    parser.add_argument('--n_epoch', default=1, type=int)
    parser.add_argument('--d_out', default=4, type=int)
    parser.add_argument('--dataset', default='thyroid', type=str)
    parser.add_argument('--exp', default='affine', type=str)
    parser.add_argument('--c_pr', default=0, type=int)
    parser.add_argument('--true_label', default=1, type=int)
    parser.add_argument('--ndf', default=8, type=int)
    parser.add_argument('--m', default=1, type=float)
    parser.add_argument('--lmbda', default=0.1, type=float)
    parser.add_argument('--eps', default=0, type=float)
    parser.add_argument('--n_iters', default=500, type=int)

    args = parser.parse_args()
    print("Dataset: ", args.dataset)

    if args.dataset == 'thyroid' or args.dataset == 'arrhythmia':
        n_iters = args.n_iters
        f_scores = np.zeros(n_iters)
        for i in range(n_iters):
            f_scores[i] = train_anomaly_detector(args)
        print("AVG f1_score", f_scores.mean())
    else:
        train_anomaly_detector(args)

transformations.py

import abc
import itertools
import numpy as np
from keras.preprocessing.image import apply_affine_transform
# The code is adapted from https://github.com/izikgo/AnomalyDetectionTransformations/blob/master/transformations.py

def get_transformer(type_trans):
    if type_trans == 'complicated':
        tr_x, tr_y = 8, 8
        transformer = Transformer(tr_x, tr_y)
        return transformer
    elif type_trans == 'simple':
        transformer = SimpleTransformer()
        return transformer


class AffineTransformation(object):
    def __init__(self, flip, tx, ty, k_90_rotate):
        self.flip = flip
        self.tx = tx
        self.ty = ty
        self.k_90_rotate = k_90_rotate

    def __call__(self, x):
        res_x = x
        if self.flip:
            res_x = np.fliplr(res_x)
        if self.tx != 0 or self.ty != 0:
            res_x = apply_affine_transform(res_x,
            tx=self.tx, ty=self.ty, channel_axis=2, fill_mode='reflect')
        if self.k_90_rotate != 0:
            res_x = np.rot90(res_x, self.k_90_rotate)
        return res_x


class AbstractTransformer(abc.ABC):
    def __init__(self):
        self._transformation_list = None
        self._create_transformation_list()

    @property
    def n_transforms(self):
        return len(self._transformation_list)

    @abc.abstractmethod
    def _create_transformation_list(self):
        return

    def transform_batch(self, x_batch, t_inds):
        assert len(x_batch) == len(t_inds)

        transformed_batch = x_batch.copy()
        for i, t_ind in enumerate(t_inds):
            transformed_batch[i] = self._transformation_list[t_ind](transformed_batch[i])
        return transformed_batch


class Transformer(AbstractTransformer):
    def __init__(self, translation_x=8, translation_y=8):
        self.max_tx = translation_x
        self.max_ty = translation_y
        super().__init__()

    def _create_transformation_list(self):
        transformation_list = []
        for is_flip, tx, ty, k_rotate in itertools.product((False, True),
                                                           (0, -self.max_tx, self.max_tx),
                                                           (0, -self.max_ty, self.max_ty),
                                                           range(4)):
            transformation = AffineTransformation(is_flip, tx, ty, k_rotate)
            transformation_list.append(transformation)
        self._transformation_list = transformation_list
        return transformation_list


class SimpleTransformer(AbstractTransformer):
    def _create_transformation_list(self):
        transformation_list = []
        for is_flip, k_rotate in itertools.product((False, True),
                                                    range(4)):
            transformation = AffineTransformation(is_flip, 0, 0, k_rotate)
            transformation_list.append(transformation)
        self._transformation_list = transformation_list
        return transformation_list

wideresnet.py

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# The code is adapted from https://github.com/xternalz/WideResNet-pytorch/blob/master/wideresnet.py


class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, dropRate=0.0):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.droprate = dropRate
        self.equalInOut = (in_planes == out_planes)
        self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
                               padding=0, bias=False) or None

    def forward(self, x):
        if not self.equalInOut:
            x = self.relu1(self.bn1(x))
        else:
            out = self.relu1(self.bn1(x))
        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
        out = self.conv1(out if self.equalInOut else x)  # todo

        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        return torch.add(x if self.equalInOut else self.convShortcut(x), out)


class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate)

    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate):
        layers = []
        for i in range(int(nb_layers)):
            layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate))
        return nn.Sequential(*layers)

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


class WideResNet(nn.Module):
    def __init__(self, depth, num_classes, widen_factor=1, dropRate=0.0):
        super(WideResNet, self).__init__()
        nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor]
        # nChannels = [16, 16*widen_factor, 32*widen_factor, 32*widen_factor]

        assert((depth - 4) % 6 == 0)
        n = (depth - 4) / 6
        block = BasicBlock
        # 1st conv before any network block
        self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=1,
                               padding=1, bias=False)
        # 1st block
        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[3])
        self.relu = nn.ReLU(inplace=True)
        self.fc = nn.Linear(nChannels[3], num_classes) # todo
        # self.fout = nn.Linear(ndf, num_classes)
        self.nChannels = nChannels[3]
        # self.softmax = nn.Softmax()

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.bias.data.zero_()

    def forward(self, x):
        out = self.conv1(x)
        out = self.block1(out)
        out = self.block2(out)
        out = self.block3(out)

        out = self.relu(self.bn1(out))
        out = F.avg_pool2d(out, 8)
        act = out.view(-1, self.nChannels)
        fs = self.fc(act)
        # out = F.log_softmax(out, dim=1)
        return act, fs