Exemplo n.º 1
0
    def build(self, x, y, weight):
        with scope("x"):
            x = placeholder(tf.float32, [None, self.dim_x], x, "x")
        with scope("y"):
            y = placeholder(tf.float32, [None], y, "y")

        gx = self.gen(x)
        dx, dgx = self.dis(x), self.dis(gx)

        with scope("loss"):
            d_loss_real = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.ones_like(dx) * 0.9, logits=dx))
            d_loss_fake = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.zeros_like(dgx), logits=dgx))
            d_loss = d_loss_real + d_loss_fake

            epsilon = 1e-10
            loss_rec = tf.reduce_mean(
                -tf.reduce_sum(x * tf.log(epsilon + gx) +
                               (1 - x) * tf.log(epsilon + 1 - gx),
                               axis=1))
            g_loss = weight* loss_rec \
                +  tf.reduce_mean(
                    tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(dgx), logits=dgx))

        with scope("AUC"):
            _, auc_dgx = tf.metrics.auc(y, tf.nn.sigmoid(dgx))
            _, auc_dx = tf.metrics.auc(y, tf.nn.sigmoid(dx))
            _, auc_gx = tf.metrics.auc(y, tf.reduce_mean((x - gx)**2, axis=1))

        g_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                   scope="generator")
        d_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                   scope="discriminator")

        with scope('train_step'):
            step = tf.train.get_or_create_global_step()
            optimizer = tf.train.AdamOptimizer()
            d_step = optimizer.minimize(d_loss, step, var_list=d_vars)
            g_step = optimizer.minimize(g_loss, step, var_list=g_vars)

        return AEGAN(self,
                     step=step,
                     x=x,
                     y=y,
                     gx=gx,
                     auc_dgx=auc_dgx,
                     auc_gx=auc_gx,
                     auc_dx=auc_dx,
                     g_step=g_step,
                     d_step=d_step,
                     g_loss=g_loss,
                     d_loss=d_loss)
Exemplo n.º 2
0
Arquivo: dae.py Projeto: jmilde/gan
    def build(self, x, y, lr_max, mult):
        with tf.variable_scope("x"):
            x = placeholder(tf.float32, [None, self.dim_x], x, "x")
        with tf.variable_scope("y"):
            y = placeholder(tf.float32, [None], y, "y")

        gx = self.gen(x)
        dx, dgx = self.dis(x), self.dis(gx)

        with tf.variable_scope("loss"):
            a = tf.reduce_mean(tf.abs(x - dx))
            b = tf.reduce_mean(tf.abs(gx - dgx))
            c = tf.reduce_mean(tf.abs(x - gx))
            d_vs_g = a - (b + c) / 2  # for balancing the learnign rate
            lr_d = sigmoid(d_vs_g, mult=mult)
            lr_g = (tf.constant(1.0) - lr_d) * lr_max
            lr_d = lr_d * lr_max

            # balance parameter for discriminator caring more about autoencoding real, or discriminating fake
            sigma = 0.5
            w_fake = tf.clip_by_value(
                sigmoid(b * sigma - a, shift=0., mult=mult), 0., 0.9
            )  # hold the discrim proportion fake aways at less than half
            d_loss = a - b * w_fake

            # weights for generator
            wg_fake = tf.clip_by_value(sigmoid(b - c, shift=0., mult=mult), 0.,
                                       1.0)
            wg_reconstruct = 1 - wg_fake
            g_loss = b * wg_fake + c * wg_reconstruct

        with tf.variable_scope("AUC"):
            _, auc_dgx = tf.metrics.auc(y, tf.reduce_mean((x - dgx)**2,
                                                          axis=1))
            _, auc_dx = tf.metrics.auc(y, tf.reduce_mean((x - dx)**2, axis=1))
            _, auc_gx = tf.metrics.auc(y, tf.reduce_mean((x - gx)**2, axis=1))

        with scope('down'):
            g_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                       scope="generator")
            d_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                       scope="discriminator")
            step = tf.train.get_or_create_global_step()
            d_step = tf.train.AdamOptimizer(lr_d).minimize(d_loss,
                                                           step,
                                                           var_list=d_vars)
            g_step = tf.train.AdamOptimizer(lr_g).minimize(g_loss,
                                                           step,
                                                           var_list=g_vars)

        return DAE(self,
                   step=step,
                   x=x,
                   y=y,
                   gx=gx,
                   dgx=dgx,
                   dx=dx,
                   auc_dgx=auc_dgx,
                   auc_gx=auc_gx,
                   auc_dx=auc_dx,
                   g_loss=g_loss,
                   d_loss=d_loss,
                   d_step=d_step,
                   g_step=g_step)
Exemplo n.º 3
0
    def build(self, x, y, context_weight, loss, lam=0., weight_type="normal"):
        with scope("x"):
            x = placeholder(tf.float32, [None, None, None, self.channel_x], x,
                            "x")
        with scope("y"):
            y = placeholder(tf.float32, [None], y, "y")

        gx = self.gen(x)
        dx = {k: v(x) for k, v in self.dis.items()}
        dgx = {k: v(gx) for k, v in self.dis.items()}
        #dx, dgx = self.dis(x), self.dis(gx)

        with scope("loss"):
            d_loss = [tf.reduce_mean(
                          tf.nn.sigmoid_cross_entropy_with_logits(
                              labels=tf.ones_like(dx[k])*0.9, logits=dx[k])) \
                      + \
                      tf.reduce_mean(
                          tf.nn.sigmoid_cross_entropy_with_logits(
                              labels=tf.zeros_like(dgx[k]), logits=dgx[k]))
                      for k in dx.keys()]

            ### old d_loss
            #d_loss_real, d_loss_fake = [], []
            #for k in dx.keys():
            #d_loss_real.append(tf.reduce_mean(
            #tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(dx[k])*0.9, logits=dx[k])))
            #d_loss_fake.append(tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(dgx[k]), logits=dgx[k])))

            #if loss=="mean":
            #d_loss = tf.reduce_mean(d_loss_real) + tf.reduce_mean(d_loss_fake)
            #elif loss=="max":
            #d_loss = tf.reduce_mean(d_loss_real) + tf.reduce_mean(d_loss_fake)
            #elif loss=="softmax":
            #d_loss = tf.reduce_mean(d_loss_real) + tf.reduce_mean(d_loss_fake)

            epsilon = 1e-10
            loss_rec = tf.reduce_mean(
                -tf.reduce_sum(x * tf.log(epsilon + gx) +
                               (1 - x) * tf.log(epsilon + 1 - gx),
                               axis=1))
            loss_g_fake = [
                tf.reduce_mean(
                    tf.nn.sigmoid_cross_entropy_with_logits(
                        labels=tf.ones_like(dgx_), logits=dgx_))
                for dgx_ in dgx.values()
            ]

            lam = placeholder(tf.float32, None, lam,
                              "lam")  # only for softmax, otherwise dummy

            with scope("lambda"):
                if loss == "softmax_self_challenged":
                    trained_l = tf.Variable(initial_value=-2.,
                                            name='controlled_lambda')
                    used_lam = tf.nn.softplus(trained_l, name='used_lambda')
                else:
                    used_lam = lam

            if loss == "mean":
                gl_adv = tf.reduce_mean(loss_g_fake)
                g_loss = context_weight * loss_rec + gl_adv
            elif loss == "max":  # max picks biggest loss = best discriminators feedback is used
                gl_adv = tf.reduce_max(loss_g_fake)
                g_loss = context_weight * loss_rec + gl_adv

            elif "softmax" in loss:
                # if lambda is self_learnt
                if used_lam == 0.:
                    weights = tf.ones_like(loss_g_fake)
                else:
                    if weight_type == 'log':
                        weights = tf.pow(loss_g_fake, used_lam)
                    else:
                        weights = tf.exp(used_lam * loss_g_fake)

                gl_adv = weighted_arithmetic(weights, loss_g_fake)

                if loss == "softmax":
                    g_loss = context_weight * loss_rec + gl_adv
                else:
                    g_loss = context_weight * loss_rec + gl_adv - 0.001 * used_lam
                #g_loss = weight* loss_rec + tf.reduce_mean(tf.nn.softmax(loss_g_fake)*loss_g_fake)

        with scope("AUC"):
            #_, auc_dgx = tf.metrics.auc(y, tf.nn.sigmoid(tf.reduce_mean(list(dgx.values()))))
            #_, auc_dx = tf.metrics.auc(y, tf.nn.sigmoid(tf.reduce_mean(list(dx.values()))))
            _, auc_gx = tf.metrics.auc(
                y, tf.reduce_mean((x - gx)**2, axis=(1, 2, 3)))

        g_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                   scope="generator")
        if loss == "softmax_self_challenged":
            lambda_var = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                           scope="loss/lambda")
            g_vars.extend(lambda_var)
        #d_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="discriminator")
        d_vars = {
            i: tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                 scope=f"discriminator_{i}")
            for i in dx.keys()
        }

        with scope('train_step'):
            step = tf.train.get_or_create_global_step()
            optimizer = tf.train.AdamOptimizer()
            #d_step = optimizer.minimize(d_loss, step, var_list=d_vars)
            d_step = [
                optimizer.minimize(loss, var_list=d_vars[i])
                for i, loss in enumerate(d_loss)
            ]
            g_step = optimizer.minimize(g_loss, step, var_list=g_vars)

        return MG_GAN(
            self,
            lam=used_lam,
            step=step,
            x=x,
            y=y,
            gx=gx
            #, auc_dgx=auc_dgx
            ,
            auc_gx=auc_gx
            #, auc_dx=auc_dx
            ,
            g_step=g_step,
            d_step=d_step,
            gl_rec=context_weight * loss_rec,
            gl_lam=0.001 * used_lam,
            gl_adv=gl_adv,
            g_loss=g_loss,
            d_loss=d_loss,
            d_loss_mean=tf.reduce_mean(d_loss),
            d_max=tf.argmax(d_loss))
Exemplo n.º 4
0
def train(anomaly_class=8):
    #set gpu
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    #load data
    (train_images,
     train_labels), (test_images,
                     test_labels) = tf.keras.datasets.mnist.load_data()
    inlier = train_images[train_labels != anomaly_class]
    x_train = np.reshape(inlier, (len(inlier), 28 * 28)) / 255
    #y_train = train_labels[train_labels!=anomaly_class]
    y_train = np.zeros(len(x_train), dtype=np.int8)  # dummy

    outlier = train_images[train_labels == anomaly_class]
    x_test = np.reshape(np.concatenate([outlier, test_images]),
                        (len(outlier) + len(test_images), 28 * 28)) / 255
    y_test = np.concatenate(
        [train_labels[train_labels == anomaly_class], test_labels])
    y_test = [0 if y != anomaly_class else 1 for y in y_test]
    x_test, y_test = unison_shfl(x_test, np.array(y_test))

    path_log = "/cache/tensorboard-logdir/ae"
    path_ckpt = "/project/outlier_detection/ckpt"

    epochs = 400
    batch_size = 700
    dim_btlnk = 32
    mult = 20
    lr_max = 1e-4
    trial = f"dae{anomaly_class}_b{batch_size}_btlnk{dim_btlnk}_lr_{lr_max}m{mult}"
    #trial="test1"
    dim_x = len(x_train[0])

    #reset graphs and fix seeds
    tf.reset_default_graph()
    if 'sess' in globals(): sess.close()
    rand = RandomState(0)
    tf.set_random_seed(0)

    # data pipeline
    batch_fn = lambda: batch2(x_train, y_train, batch_size)
    x, y = pipe(batch_fn, (tf.float32, tf.float32), prefetch=4)
    #z = tf.random_normal((batch_size, z_dim))

    # load graph
    dae = DAE.new(dim_x, dim_btlnk)
    model = DAE.build(dae, x, y, lr_max, mult)

    # start session, initialize variables
    sess = tf.InteractiveSession()
    saver = tf.train.Saver()

    wrtr = tf.summary.FileWriter(pform(path_log, trial))
    #wrtr.add_graph(sess.graph)

    ### if load pretrained model
    # pretrain = "modelname"
    #saver.restore(sess, pform(path_ckpt, pretrain))
    ### else:
    auc_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope='AUC')
    init = tf.group(tf.global_variables_initializer(),
                    tf.variables_initializer(var_list=auc_vars))
    sess.run(init)

    def log(step,
            wrtr=wrtr,
            log=tf.summary.merge([
                tf.summary.scalar('g_loss', model.g_loss),
                tf.summary.scalar('d_loss', model.d_loss),
                tf.summary.image('gx400',
                                 spread_image(model.gx[:400], 20, 20, 28, 28)),
                tf.summary.image('dgx400',
                                 spread_image(model.dgx[:400], 20, 20, 28,
                                              28)),
                tf.summary.image('dx400',
                                 spread_image(model.dx[:400], 20, 20, 28, 28)),
                tf.summary.scalar("AUC_dgx", model.auc_dgx),
                tf.summary.scalar("AUC_dx", model.auc_dx),
                tf.summary.scalar("AUC_gx", model.auc_gx)
            ]),
            y=y_test,
            x=x_test):
        wrtr.add_summary(sess.run(log, {model.x: x, model.y: y}), step)
        wrtr.flush()

    steps_per_epoch = len(x_train) // batch_size
    for epoch in tqdm(range(epochs)):
        for i in range(steps_per_epoch):
            sess.run(model.d_step)
            sess.run(model.g_step)

        # tensorboard writer
        log(sess.run(model["step"]) // steps_per_epoch)

    saver.save(sess, pform(path_ckpt, trial), write_meta_graph=False)
Exemplo n.º 5
0
def train(anomaly_class, loss_type):
    #set gpu
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"

    #load data
    (train_images,
     train_labels), (test_images,
                     test_labels) = tf.keras.datasets.mnist.load_data()
    inlier = train_images[train_labels != anomaly_class]
    x_train = np.reshape(inlier, (len(inlier), 28 * 28)) / 255
    #y_train = train_labels[train_labels!=anomaly_class]
    y_train = np.zeros(len(x_train), dtype=np.int8)  # dummy
    outlier = train_images[train_labels == anomaly_class]
    x_test = np.reshape(np.concatenate([outlier, test_images]),
                        (len(outlier) + len(test_images), 28 * 28)) / 255
    y_test = np.concatenate(
        [train_labels[train_labels == anomaly_class], test_labels])
    y_test = [0 if y != anomaly_class else 1 for y in y_test]
    x_test, y_test = unison_shfl(x_test, np.array(y_test))

    path_log = "/cache/tensorboard-logdir/ae"
    path_ckpt = "/project/outlier_detection/ckpt"

    epochs = 400
    batch_size = 700
    dim_btlnk = 32
    dim_z = dim_btlnk
    dim_dense = 32
    accelerate = 1e-5
    context_weight = 1
    trial = f"vaegan_{loss_type}_{anomaly_class}_b{batch_size}_btlnk{dim_btlnk}_d{dim_dense}_n{dim_z}_a{accelerate}"

    dim_x = len(x_train[0])
    #reset graphs and fix seeds
    tf.reset_default_graph()
    if 'sess' in globals(): sess.close()
    rand = RandomState(0)
    tf.set_random_seed(0)

    # data pipeline
    batch_fn = lambda: batch2(x_train, y_train, batch_size, dim_z)
    x, y, z = pipe(batch_fn, (tf.float32, tf.float32, tf.float32), prefetch=4)

    # load graph
    aegan = VAEGAN.new(dim_x, dim_btlnk, dim_dense, dim_z, accelerate)
    model = VAEGAN.build(aegan, x, y, z, loss_type)

    # start session, initialize variables

    sess = tf.InteractiveSession()
    saver = tf.train.Saver()

    wrtr = tf.summary.FileWriter(pform(path_log, trial))
    #wrtr.add_graph(sess.graph)

    ### if load pretrained model
    # pretrain = "modelname"
    #saver.restore(sess, pform(path_ckpt, pretrain))
    ### else:
    auc_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope='AUC')
    init = tf.group(tf.global_variables_initializer(),
                    tf.variables_initializer(var_list=auc_vars))
    sess.run(init)

    def log(
            step,
            wrtr=wrtr,
            log=tf.summary.
        merge([
            tf.summary.scalar('g_loss', model.g_loss),
            tf.summary.scalar('d_loss', model.d_loss),
            tf.summary.scalar('mu', model.m),
            tf.summary.scalar('lv', model.l),
            tf.summary.image('gzx400',
                             spread_image(model.gzx[:400], 20, 20, 28, 28))
            #, tf.summary.image('gz400', spread_image(model.gz[:400], 20,20,28,28))
            ,
            tf.summary.scalar("AUC_gzx", model.auc_gzx),
            tf.summary.scalar("AUC_dgzx", model.auc_dgzx),
            tf.summary.scalar("AUC_dx", model.auc_dx)
            #, tf.summary.scalar("gz_loss",model.gz_loss)
            ,
            tf.summary.scalar("gzx_loss", model.gzx_loss),
            tf.summary.scalar("ftr_loss", model.ftr_loss),
            tf.summary.scalar("kl_loss", model.kl_loss),
            tf.summary.scalar("dx_loss", model.dx_loss)
            #, tf.summary.scalar("dgz_loss",model.dgz_loss)
            ,
            tf.summary.scalar("dgzx_loss", model.dgzx_loss)
        ]),
            y=y_test,
            x=x_test):
        mu = sess.run(model.mu, {model.x: x})
        wrtr.add_summary(sess.run(log, {
            model.zx: mu,
            model.x: x,
            model.y: y
        }), step)
        wrtr.flush()

    steps_per_epoch = len(x_train) // batch_size
    for epoch in tqdm(range(epochs)):
        for i in range(steps_per_epoch):
            sess.run(model.g_step)
            sess.run(model.d_step)

        # tensorboard writer
        #log(sess.run(model["step"])//steps_per_epoch)
        log(sess.run(model["step"] // steps_per_epoch))

    saver.save(sess, pform(path_ckpt, trial), write_meta_graph=False)
Exemplo n.º 6
0
model.lr   = train[0].lr
model.step = train[0].step
model.errt = train[0].errt
model.loss = tf.add_n([t.loss for t in train])
model.down = tf.train.AdamOptimizer(model.lr, T.beta1, T.beta2, T.epsilon) \
                     .minimize(model.loss, model.step, (model.embeds[1].logit, model.embeds[3].logit))

############
# training #
############

sess = tf.InteractiveSession()

tf.global_variables_initializer().run()
tf.train.Saver(
    [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if not ('embed' in v.name and 'Adam' in v.name)]
).restore(sess, pform(P.ckpt, 'm1_', 9))

saver = tf.train.Saver()

model.step.assign(0).eval()

def summ(step, wtr = tf.summary.FileWriter(pform(P.log, C.trial))
         , summary = tf.summary.merge(
             ( tf.summary.scalar('step_errt', model.errt)
             , tf.summary.scalar('step_loss', model.loss)))):
    errt, loss = map(comp(np.mean, np.concatenate), zip(*chain(*(
        batch_run(sess, m, (m.errt_samp, m.loss_samp), s, t, batch= C.batch_valid)
        for m, (s, t) in zip(valid, data_valid)))))
    wtr.add_summary(sess.run(summary, {model.errt: errt, model.loss: loss}), step)
    wtr.flush()
Exemplo n.º 7
0
    def build(self, x, y, z, loss_type):
        d_scale_factor = tf.constant(0.)  #tf.constant(0.25)
        g_scale_factor = tf.constant(0.)  #tf.constant(1 - 0.75/2)
        with scope("x"):
            x = placeholder(tf.float32, [None, self.dim_x], x, "x")
        with scope("y"):
            y = placeholder(tf.float32, [None], y, "y")
        with scope("z"):
            z = placeholder(tf.float32, [None, self.dim_noise], z, "z")

        zx, mu, lv, hl_e = self.enc(x)

        gzx = self.gen(zx)
        #gz = self.gen(z)

        dx, hl_dx = self.dis(x)
        dgzx, hl_dgzx = self.dis(gzx)
        #dgz, hl_dgz = self.dis(gz)

        with tf.variable_scope("step"):
            step = tf.train.get_or_create_global_step()
            rate = self.accelerate * tf.to_float(step)
            rate_anneal = tf.tanh(rate)

        with scope("loss"):
            dx_loss = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.ones_like(dx) - d_scale_factor, logits=dx))
            dgzx_loss = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.zeros_like(dgzx), logits=dgzx))
            #dgz_loss =  tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(dgz), logits=dgz))
            d_loss = dx_loss + dgzx_loss  #+ dgz_loss

            kl_loss = tf.reduce_mean(0.5 *
                                     (tf.square(mu) + tf.exp(lv) - lv - 1.0))

            gzx_loss = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    labels=tf.ones_like(dgzx) - g_scale_factor, logits=dgzx))

            if loss_type == "xtrpy":
                epsilon = 1e-10
                ftr_loss = tf.reduce_mean(
                    -tf.reduce_sum(x * tf.log(epsilon + gzx) +
                                   (1 - x) * tf.log(epsilon + 1 - gzx),
                                   axis=1))
                g_loss = gzx_loss / 10 + ftr_loss / 5 + kl_loss * rate_anneal

            else:
                ftr_loss = tf.reduce_mean(tf.abs(x - gzx))
                g_loss = gzx_loss / 2 + ftr_loss * 10 + kl_loss * rate_anneal

        with scope("AUC"):
            _, auc_gzx = tf.metrics.auc(y, tf.reduce_mean((x - gzx)**2,
                                                          axis=1))
            _, auc_dx = tf.metrics.auc(y, tf.nn.sigmoid(dx))
            _, auc_dgzx = tf.metrics.auc(y, tf.nn.sigmoid(dgzx))

        g_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                   scope="generator")
        g_vars.append(
            tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope="encoder"))
        print(g_vars)
        d_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                   scope="discriminator")
        print(d_vars)

        with scope('train_step'):
            #optimizer = tf.train.RMSPropOptimizer()
            optimizer = tf.train.AdamOptimizer()
            d_step = optimizer.minimize(d_loss, step, var_list=d_vars)
            g_step = optimizer.minimize(g_loss, step, var_list=g_vars)

        return VAEGAN(
            self,
            step=step,
            x=x,
            y=y,
            z=z,
            zx=zx,
            mu=mu,
            lv=lv,
            m=tf.reduce_mean(mu),
            l=tf.reduce_mean(lv)
            #, gz=gz
            ,
            gzx=gzx,
            auc_gzx=auc_gzx,
            auc_dx=auc_dx,
            auc_dgzx=auc_dgzx,
            g_step=g_step,
            d_step=d_step,
            g_loss=g_loss,
            d_loss=d_loss
            #,gz_loss=gz_loss
            ,
            gzx_loss=gzx_loss,
            ftr_loss=ftr_loss,
            kl_loss=kl_loss,
            dx_loss=dx_loss
            #, dgz_loss=dgz_loss
            ,
            dgzx_loss=dgzx_loss)
def train(anomaly_class=8,
          dataset="cifar",
          n_dis=1,
          epochs=25,
          dim_btlnk=32,
          batch_size=64,
          loss="mean",
          context_weight=1,
          dim_d=64,
          dim_g=64,
          extra_layers=0,
          gpu="0"):

    #set gpu
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu

    path_log = f"/cache/tensorboard-logdir/{dataset}"
    path_ckpt = "/project/multi-discriminator-gan/ckpt"
    path_data = "/project/multi-discriminator-gan/data"

    #reset graphs and fix seeds
    tf.reset_default_graph()
    if 'sess' in globals(): sess.close()
    rand = RandomState(0)
    tf.set_random_seed(0)

    #load data
    if dataset == "ucsd1":
        x_train = np.load("./data/ucsd1_train_x.npz")["arr_0"] / 255
        y_train = np.load("./data/ucsd1_train_y.npz")["arr_0"]
        x_test = np.load("./data/ucsd1_test_x.npz")["arr_0"] / 255
        y_test = np.load("./data/ucsd1_test_y.npz")["arr_0"]

    elif dataset == "uscd2":
        x_train = np.load("./data/ucsd2_train_x.npz")["arr_0"]
        y_train = np.load("./data/ucsd2_train_y.npz")["arr_0"]
        x_test = np.load("./data/ucsd2_test_x.npz")["arr_0"]
        y_test = np.load("./data/ucsd2_test_y.npz")["arr_0"]

    else:
        if dataset == "mnist":
            (train_images, train_labels), (
                test_images,
                test_labels) = tf.keras.datasets.mnist.load_data()
            train_images = resize_images(train_images)
            test_images = resize_images(test_images)
        else:
            (train_images, train_labels), (
                test_images,
                test_labels) = tf.keras.datasets.cifar10.load_data()
            train_labels = np.reshape(train_labels, len(train_labels))
            test_labels = np.reshape(test_labels, len(test_labels))

        inlier = train_images[train_labels != anomaly_class]
        #data_size = prod(inlier[0].sha
        x_train = inlier / 255
        #x_train = np.reshape(inlier, (len(inlier), data_size))/255
        #y_train = train_labels[train_labels!=anomaly_class]
        y_train = np.zeros(len(x_train), dtype=np.int8)  # dummy
        outlier = train_images[train_labels == anomaly_class]
        x_test = np.concatenate([outlier, test_images]) / 255
        #x_test = np.reshape(np.concatenate([outlier, test_images])
        #                    ,(len(outlier)+len(test_images), data_size))/255
        y_test = np.concatenate(
            [train_labels[train_labels == anomaly_class], test_labels])
        y_test = [0 if y != anomaly_class else 1 for y in y_test]
        x_test, y_test = unison_shfl(x_test, np.array(y_test))

    img_size_x = x_train[0].shape[0]
    img_size_y = x_train[0].shape[1]
    channel = x_train[0].shape[-1]
    trial = f"{dataset}_{loss}_dis{n_dis}_{anomaly_class}_w{context_weight}_btlnk{dim_btlnk}_d{dim_d}_g{dim_g}e{extra_layers}"

    # data pipeline
    batch_fn = lambda: batch2(x_train, y_train, batch_size)
    x, y = pipe(batch_fn, (tf.float32, tf.float32), prefetch=4)
    #z = tf.random_normal((batch_size, z_dim))

    # load graph
    mg_gan = MG_GAN.new(img_size_x,
                        channel,
                        dim_btlnk,
                        dim_d,
                        dim_g,
                        n_dis,
                        extra_layers=0)
    model = MG_GAN.build(mg_gan, x, y, context_weight, loss)

    # start session, initialize variables

    sess = tf.InteractiveSession()
    saver = tf.train.Saver()

    wrtr = tf.summary.FileWriter(pform(path_log, trial))
    wrtr.add_graph(sess.graph)

    ### if load pretrained model
    # pretrain = "modelname"
    #saver.restore(sess, pform(path_ckpt, pretrain))
    ### else:
    auc_vars = tf.get_collection(tf.GraphKeys.LOCAL_VARIABLES, scope='AUC')
    init = tf.group(tf.global_variables_initializer(),
                    tf.variables_initializer(var_list=auc_vars))
    sess.run(init)

    #if "ucsd" in dataset:
    summary_test = tf.summary.merge([
        tf.summary.scalar('g_loss', model.g_loss),
        tf.summary.scalar("lambda", model.lam),
        tf.summary.scalar("gl_rec", model.gl_rec),
        tf.summary.scalar("gl_adv", model.gl_adv),
        tf.summary.scalar("gl_lam", model.gl_lam),
        tf.summary.scalar('d_loss_mean', model.d_loss_mean),
        tf.summary.scalar('d_max', model.d_max)
        #, tf.summary.scalar('d_loss', model.d_loss)
        ,
        tf.summary.scalar("AUC_gx", model.auc_gx)
    ])
    if dataset == "ucsd1":
        summary_images = tf.summary.merge(
            (tf.summary.image("gx", model.gx, max_outputs=8),
             tf.summary.image("x", model.x, max_outputs=8),
             tf.summary.image(
                 'gx400',
                 spread_image(tf.concat([model.gx, model.x], axis=1), 8, 2,
                              img_size_x, img_size_y, channel))))
    else:
        summary_images = tf.summary.merge(
            (tf.summary.image("gx", model.gx, max_outputs=8),
             tf.summary.image(
                 'gx400',
                 spread_image(model.gx[:400], 20, 20, img_size_x, img_size_y,
                              channel)),
             tf.summary.image("x", model.x, max_outputs=8)))

    if n_dis > 1:
        d_wrtr = {
            i: tf.summary.FileWriter(pform(path_log, trial + f"d{i}"))
            for i in range(n_dis)
        }
        summary_discr = {
            i: tf.summary.scalar('d_loss_multi', model.d_loss[i])
            for i in range(n_dis)
        }

    def summ(step):
        fetches = model.g_loss, model.lam, model.d_loss_mean, model.auc_gx
        results = map(
            np.mean,
            zip(*(sess.run(fetches, {
                model['x']: x_test[i:j],
                model['y']: y_test[i:j]
            }) for i, j in partition(len(x_test), batch_size, discard=False))))
        results = list(results)
        wrtr.add_summary(sess.run(summary_test, dict(zip(fetches, results))),
                         step)

        if dataset == "ucsd1":
            # bike, skateboard, grasswalk, shopping cart, car, normal, normal, grass
            wrtr.add_summary(
                sess.run(
                    summary_images, {
                        model.x:
                        x_test[[990, 1851, 2140, 2500, 2780, 2880, 3380, 3580]]
                    }), step)
        else:
            wrtr.add_summary(sess.run(summary_images, {model.x: x_test}), step)
        wrtr.flush()

    def summ_discr(step):
        fetches = model.d_loss
        results = map(
            np.mean,
            zip(*(sess.run(fetches, {
                model['x']: x_test[i:j],
                model['y']: y_test[i:j]
            }) for i, j in partition(len(x_test), batch_size, discard=False))))
        results = list(results)
        if n_dis > 1:  # put all losses of the discriminators in one plot
            for i in range(n_dis):
                d_wrtr[i].add_summary(
                    sess.run(summary_discr[i], dict(zip(fetches, results))),
                    step)
                #d_wrtr[i].add_summary(sess.run(summary_discr[i], dict([(fetches[i], results[i])])), step)
                d_wrtr[i].flush()

    #def log(step
    #        , wrtr= wrtr
    #        , log = tf.summary.merge([tf.summary.scalar('g_loss', model.g_loss)
    #                                  , tf.summary.scalar('d_loss', tf.reduce_mean(model.d_loss))
    #                                  , tf.summary.scalar("lambda", model.lam)
    #                                  , tf.summary.image("gx", model.gx, max_outputs=5)
    #                                  , tf.summary.image('gx400', spread_image(model.gx[:400], 20,20, img_size, img_size, channel))
    #                                  #, tf.summary.scalar("AUC_dgx", model.auc_dgx)
    #                                  #, tf.summary.scalar("AUC_dx", model.auc_dx)
    #                                  , tf.summary.scalar("AUC_gx", model.auc_gx)])
    #        , y= y_test
    #        , x= x_test):
    #    wrtr.add_summary(sess.run(log, {model["x"]:x
    #                                    , model["y"]:y})
    #                     , step)
    #    wrtr.flush()

    steps_per_epoch = len(x_train) // batch_size - 1
    for epoch in tqdm(range(epochs)):
        for i in range(steps_per_epoch):
            #sess.run(model["train_step"])
            sess.run(model['d_step'])
            sess.run(model['g_step'])
        # tensorboard writer
        #if "ucsd" in dataset:
        summ(sess.run(model["step"]) // steps_per_epoch)
        #else:
        #    log(sess.run(model["step"])//steps_per_epoch)
        if n_dis > 1:
            summ_discr(sess.run(model["step"]) // steps_per_epoch)

    saver.save(sess, pform(path_ckpt, trial), write_meta_graph=False)
Exemplo n.º 9
0
model.step = train[0].step
model.errt = train[0].errt
model.loss = tf.add_n([t.loss for t in train])
model.down = tf.train.AdamOptimizer(model.lr, T.beta1, T.beta2,
                                    T.epsilon).minimize(
                                        model.loss, model.step)

############
# training #
############

sess = tf.InteractiveSession()

tf.global_variables_initializer().run()
tf.train.Saver([
    v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
    if not ('embed' in v.name and 'Adam' in v.name)
]).restore(sess, pform(P.ckpt, 'm1_', 9))

saver = tf.train.Saver()

model.step.assign(0).eval()


def summ(step,
         wtr=tf.summary.FileWriter(pform(P.log, C.trial)),
         summary=tf.summary.merge((tf.summary.scalar('step_errt', model.errt),
                                   tf.summary.scalar('step_loss',
                                                     model.loss)))):
    errt, loss = map(
        comp(np.mean, np.concatenate),