예제 #1
0
파일: wgan_gp.py 프로젝트: exaV/molanet
    def create_generator_loss(
        self,
        x: tf.Tensor,
        y: tf.Tensor,
        generator: tf.Tensor,
        generator_discriminator: tf.Tensor,
        apply_summary: bool = True,
        use_gpu: bool = True,
        data_format: str = "NHWC"
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, List[tf.Tensor]]]:
        summary_ops = []

        with tf.device(select_device(use_gpu)):
            loss = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=generator_discriminator,
                    labels=tf.ones_like(generator_discriminator)))

        if self._l1_lambda > 0.0:
            _log.info(f"Using l1 loss, lambda={self._l1_lambda}")

            with tf.device(select_device(use_gpu)):
                l1_loss = tf.reduce_mean(tf.abs(tf.subtract(y, generator)))

            with use_cpu():
                if apply_summary:
                    summary_ops.append(
                        tf.summary.scalar("generator_loss_l1", l1_loss))
                    summary_ops.append(
                        tf.summary.scalar("generator_loss_discriminator",
                                          loss))

                loss = loss + tf.constant(self._l1_lambda,
                                          dtype=tf.float32) * l1_loss

        if apply_summary:
            summary_ops.append(tf.summary.scalar("generator_loss", loss))
            return loss, summary_ops
        else:
            return loss
예제 #2
0
    def create_generator_loss(
        self,
        x: tf.Tensor,
        y: tf.Tensor,
        generator: tf.Tensor,
        generator_discriminator: tf.Tensor,
        apply_summary: bool = True,
        use_gpu: bool = True,
        data_format: str = "NHWC"
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, List[tf.Tensor]]]:

        with tf.device(select_device(use_gpu)):
            loss = tf.reduce_mean(
                # Generator should have 0 loss
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=generator_discriminator,
                    labels=tf.zeros_like(generator_discriminator)))

        if apply_summary:
            return loss, [tf.summary.scalar("generator_loss", loss)]
        else:
            return loss
예제 #3
0
    def create_discriminator_loss(
        self,
        x: tf.Tensor,
        y: tf.Tensor,
        generator: tf.Tensor,
        generator_discriminator: tf.Tensor,
        real_discriminator: tf.Tensor,
        apply_summary: bool = True,
        use_gpu: bool = True,
        data_format: str = "NHWC"
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, List[tf.Tensor]]]:
        with use_cpu():
            batch_size = tf.shape(x)[0]
            epsilons = tf.random_uniform((batch_size, ),
                                         0.0,
                                         1.0,
                                         seed=self._seed)

            gradient_input = tf.multiply(epsilons, y) + tf.multiply(
                tf.ones_like(epsilons) - epsilons, generator)

        with tf.device(select_device(use_gpu)):
            gradient_discriminator, gradient_discriminator_input = self._network_factory.create_discriminator(
                x,
                gradient_input,
                reuse=True,
                return_input_tensor=True,
                use_gpu=use_gpu,
                data_format=data_format)

            generator_discriminator = tf.reshape(generator_discriminator,
                                                 shape=[-1])

            # Generated samples should be close to their inverse jaccard index => Loss defined
            # Generator samples have to be first converted into range [0, 1]
            generated_jaccard_index = jaccard_index(
                values=tanh_to_sigmoid(generator), labels=tanh_to_sigmoid(y))
            generated_jaccard_loss = tf.ones_like(
                generated_jaccard_index) - generated_jaccard_index

            loss_generated = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=generator_discriminator,
                    labels=generated_jaccard_loss))

            # Real samples should all be 0 => No loss
            loss_real = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=real_discriminator,
                    labels=tf.zeros_like(real_discriminator)))

            gradient = tf.gradients(gradient_discriminator,
                                    gradient_discriminator_input)
            gradient_norm = tf.norm(gradient)
            gradient_penalty_raw = tf.reduce_mean(
                (gradient_norm - tf.ones_like(gradient_norm))**2)

        with use_cpu():
            gradient_penalty = tf.multiply(
                tf.constant(self._gradient_lambda, dtype=tf.float32),
                gradient_penalty_raw)
            loss = loss_generated + loss_real + gradient_penalty

        if apply_summary:
            summary_operations = [
                tf.summary.scalar("discriminator_loss_real", loss_real),
                tf.summary.scalar("discriminator_loss_generated",
                                  loss_generated),
                tf.summary.scalar("discriminator_gradient_penalty",
                                  gradient_penalty),
                tf.summary.scalar("discriminator_gradient_norm",
                                  gradient_norm),
                tf.summary.scalar("discriminator_loss", loss)
            ]

            return loss, summary_operations
        else:
            return loss
예제 #4
0
파일: wgan_gp.py 프로젝트: exaV/molanet
    def create_discriminator_loss(
        self,
        x: tf.Tensor,
        y: tf.Tensor,
        generator: tf.Tensor,
        generator_discriminator: tf.Tensor,
        real_discriminator: tf.Tensor,
        apply_summary: bool = True,
        use_gpu: bool = True,
        data_format: str = "NHWC"
    ) -> Union[tf.Tensor, Tuple[tf.Tensor, List[tf.Tensor]]]:
        with use_cpu():
            batch_size = tf.shape(x)[0]
            epsilons = tf.random_uniform((batch_size, ),
                                         0.0,
                                         1.0,
                                         seed=self._seed)

            gradient_input = tf.multiply(epsilons, y) + tf.multiply(
                tf.ones_like(epsilons) - epsilons, generator)

        with tf.device(select_device(use_gpu)):
            gradient_discriminator, gradient_discriminator_input = self._network_factory.create_discriminator(
                x,
                gradient_input,
                reuse=True,
                return_input_tensor=True,
                use_gpu=use_gpu,
                data_format=data_format)

            loss_generated = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=generator_discriminator,
                    labels=tf.zeros_like(generator_discriminator)))

            loss_real = tf.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(
                    logits=real_discriminator,
                    labels=tf.ones_like(real_discriminator)))

            gradient = tf.gradients(gradient_discriminator,
                                    gradient_discriminator_input)
            gradient_norm = tf.norm(gradient)
            gradient_penalty_raw = tf.reduce_mean(
                (gradient_norm - tf.ones_like(gradient_norm))**2)

        with use_cpu():
            gradient_penalty = tf.multiply(
                tf.constant(self._gradient_lambda, dtype=tf.float32),
                gradient_penalty_raw)
            loss = loss_generated + loss_real + gradient_penalty

        if apply_summary:
            summary_operations = [
                tf.summary.scalar("discriminator_loss_real", loss_real),
                tf.summary.scalar("discriminator_loss_generated",
                                  loss_generated),
                tf.summary.scalar("discriminator_gradient_penalty",
                                  gradient_penalty),
                tf.summary.scalar("discriminator_gradient_norm",
                                  gradient_norm),
                tf.summary.scalar("discriminator_loss", loss)
            ]

            return loss, summary_operations
        else:
            return loss
예제 #5
0
    def train(self):
        """
        Trains the model until a maximum number of iterations as specified in the training options is reached.

        This method requires this trainer to have __enter__ called previously, otherwise no session exists
        and calls to this method will fail.
        """
        if self._sess is None:
            raise RuntimeError("A running session is required to start training")

        # Start input enqueue threads.
        coord = tf.train.Coordinator()
        _log.info("Starting queue runners...")
        threads = tf.train.start_queue_runners(sess=self._sess, coord=coord)

        save_model_path = os.path.join(self._training_options.summary_directory, "model.ckpt")
        save_image_path = os.path.join(self._training_options.summary_directory, "images/")

        if not os.path.exists(save_image_path):
            os.makedirs(save_image_path)

        with tf.device(select_device(self._training_options.use_gpu)):
            init_ops = (tf.global_variables_initializer(), tf.local_variables_initializer())

        tf.get_default_graph().finalize()

        if self._restored_iteration is None:
            self._sess.run(init_ops)

        iteration = self._sess.run(self._global_step)

        _log.info("Starting training")

        try:
            while not coord.should_stop():
                current_iteration = iteration
                iteration = self._sess.run(self._step_op)

                if current_iteration % self._training_options.save_model_interval == 0:
                    self._train_saver.save(self._sess, save_model_path, global_step=self._global_step)
                    _log.info(f"Saved model from iteration {iteration}")

                # Run CV validation
                if current_iteration % self._training_options.cv_summary_interval == 0:
                    _log.info("Evaluating CV set...")

                    # TODO: Make configurable?
                    IMAGE_OUTPUT_COUNT = 5
                    image_output_iterations = min(IMAGE_OUTPUT_COUNT, self._cv_pipeline.sample_count)

                    # Generate individual summaries
                    cv_summaries = []
                    cv_images = []
                    for _ in range(image_output_iterations):
                        current_summary, current_image = self._sess.run(
                            (self._cv_summary, self._concatenated_images_op))
                        cv_summaries.append(current_summary)
                        cv_images.append(current_image)

                    for _ in range(image_output_iterations, self._cv_pipeline.sample_count):
                        cv_summaries.append(self._sess.run(self._cv_summary))

                    # Convert summaries into numpy array and calculate the average for all tags
                    # TODO: Validate that only scalar summaries
                    summary_tags = [entry.tag for entry in tf.Summary.FromString(cv_summaries[0]).value]
                    tag_index_mapping = {tag: idx for idx, tag in enumerate(summary_tags)}
                    summary_values = np.zeros((len(cv_summaries), len(summary_tags)), dtype=np.float32)
                    for idx, current_summary in enumerate(cv_summaries):
                        for entry in tf.Summary.FromString(current_summary).value:
                            summary_values[idx, tag_index_mapping[entry.tag]] = entry.simple_value

                    real_summary = np.mean(summary_values, axis=0)

                    # Create output summary proto
                    value_list = [
                        tf.Summary.Value(tag=tag, simple_value=real_summary[idx])
                        for idx, tag in enumerate(summary_tags)]
                    result_proto = tf.Summary(value=value_list)

                    # Write summary
                    self._cv_summary_writer.add_summary(result_proto, current_iteration)

                    # Write images
                    # TODO: Don't use hardcoded size
                    for idx, current_image in enumerate(cv_images):
                        output_image = np.reshape(current_image, (512, 512 * 4, 3))
                        Image.fromarray(output_image, "RGB").save(
                            os.path.join(save_image_path, f"sample_{idx:02d}_{iteration:08d}.png"))

                # Train discriminator
                for _ in range(self._training_options.discriminator_iterations):
                    self._sess.run(self._op_discriminator)

                # Train generator, optionally output summary
                if current_iteration % self._training_options.training_summary_interval == 0:
                    _, current_summary = self._sess.run([self._op_generator, self._train_summary])

                    self._train_summary_writer.add_summary(current_summary, iteration)

                    _log.info(f"Iteration {iteration} done")
                else:
                    self._sess.run(self._op_generator)

                # Check for iteration limit reached
                if iteration > self._training_options.max_iterations:
                    coord.request_stop()

        except Exception as ex:
            coord.request_stop(ex)
        finally:
            coord.request_stop()

            _log.info("Waiting for threads to finish...")
            coord.join(threads)

            # Close writers AFTER threads stopped to make sure summaries are written
            self._train_summary_writer.close()
            self._cv_summary_writer.close()

        _log.info("Training finished")
예제 #6
0
    def __init__(
            self,
            training_pipeline: InputPipeline,
            cv_pipeline: InputPipeline,
            network_factory: NetworkFactory,
            objective_factory: ObjectiveFactory,
            training_options: TrainingOptions,
            learning_rate: float,
            beta1: float = 0.9,
            beta2: float = 0.999):
        """
        Create a new network trainer
        :param training_pipeline: Input pipeline used for training
        :param cv_pipeline: Input pipeline used for cross-validation
        :param network_factory: Factory to create training and evaluation networks
        :param objective_factory: Factory to create generator and discriminator losses
        :param training_options: Options controlling the training process
        :param learning_rate: Learning rate to use in the Adam optimizer
        :param beta1: Beta1 to use in the Adam optimizer
        :param beta2: Beta2 to use in the Adam optimizer
        """

        self._training_options = training_options
        self._restored_iteration = None

        # Create input pipelines
        with use_cpu():
            self._training_pipeline = training_pipeline
            self._train_x, self._train_y, _ = training_pipeline.create_pipeline()
            self._cv_pipeline = cv_pipeline
            self._cv_x, self._cv_y, _ = self._cv_pipeline.create_pipeline()

        # Create training graph
        with tf.name_scope("training"):

            # Create networks
            self._generator = network_factory.create_generator(self._train_x, use_gpu=self._training_options.use_gpu,
                                                               data_format=self._training_options.data_format)
            self._discriminator_generated = network_factory.create_discriminator(
                self._train_x, self._generator, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)
            self._discriminator_real = network_factory.create_discriminator(
                self._train_x, self._train_y, reuse=True, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)

            # Create losses
            self._generator_loss, generator_summary = objective_factory.create_generator_loss(
                self._train_x, self._train_y,
                self._generator, self._discriminator_generated, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)
            self._discriminator_loss, discriminator_summary = objective_factory.create_discriminator_loss(
                self._train_x, self._train_y,
                self._generator, self._discriminator_generated, self._discriminator_real,
                use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)

            with tf.device(select_device(self._training_options.use_gpu)):
                # Create optimizers
                trainable_variables = tf.trainable_variables()
                variables_discriminator = [var for var in trainable_variables if var.name.startswith("discriminator")]
                variables_generator = [var for var in trainable_variables if var.name.startswith("generator")]

                self._optimizer_generator = tf.train.AdamOptimizer(learning_rate, beta1, beta2, name="adam_generator")
                self._optimizer_discriminator = tf.train.AdamOptimizer(learning_rate, beta1, beta2, name="adam_discriminator")

                self._op_generator = self._optimizer_generator.minimize(self._generator_loss, var_list=variables_generator)
                self._op_discriminator = self._optimizer_discriminator.minimize(self._discriminator_loss, var_list=variables_discriminator)

            with use_cpu():
                # Iteration counter
                self._global_step = tf.Variable(0, trainable=False, name="global_step", dtype=tf.int64)
                self._step_op = tf.assign_add(self._global_step, 1)

            # Create summary operation
            accuracy, precision, recall, f1_score, specificity, jaccard_similarity = _create_summaries(self._generator, self._train_y)
            summary_operations = [
                tf.summary.scalar("accuracy", accuracy),
                tf.summary.scalar("precision", precision),
                tf.summary.scalar("recall", recall),
                tf.summary.scalar("f1_score", f1_score),
                tf.summary.scalar("specificity", specificity),
                tf.summary.scalar("jaccard_similarity", jaccard_similarity)
            ]

            self._train_saver = tf.train.Saver(keep_checkpoint_every_n_hours=1)

            # Merge summaries
            self._train_summary = tf.summary.merge(summary_operations + generator_summary + discriminator_summary)
            self._train_summary_writer = tf.summary.FileWriter(
                os.path.join(self._training_options.summary_directory, "training"), graph=tf.get_default_graph())

        # Create CV graph
        with tf.name_scope("cv"):
            # Create networks
            generator = network_factory.create_generator(
                self._cv_x, reuse=True, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)
            discriminator_generated = network_factory.create_discriminator(
                self._cv_x, generator, reuse=True, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)
            discriminator_real = network_factory.create_discriminator(
                self._cv_x, self._cv_y, reuse=True, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)

            # Create losses
            _, generator_summary = objective_factory.create_generator_loss(
                self._cv_x, self._cv_y, generator, discriminator_generated, use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)
            _, discriminator_summary = objective_factory.create_discriminator_loss(
                self._cv_x, self._cv_y, generator, discriminator_generated, discriminator_real,
                use_gpu=self._training_options.use_gpu,
                data_format=self._training_options.data_format)

            # Create other summary options
            accuracy, precision, recall, f1_score, specificity, jaccard_similarity = _create_summaries(generator, self._cv_y)

            # Create summary operation
            summary_operations = [
                tf.summary.scalar("accuracy", accuracy),
                tf.summary.scalar("precision", precision),
                tf.summary.scalar("recall", recall),
                tf.summary.scalar("f1_score", f1_score),
                tf.summary.scalar("specificity", specificity),
                tf.summary.scalar("jaccard_similarity", jaccard_similarity)
            ]

            with use_cpu():
                # Concatenated images
                self._concatenated_images_op = _create_concatenated_images(
                    self._cv_x,
                    self._cv_y,
                    generator,
                    self._cv_pipeline.color_converter,
                    self._training_options.data_format
                )

            # Merge summaries
            self._cv_summary = tf.summary.merge(summary_operations + generator_summary + discriminator_summary)
            self._cv_summary_writer = tf.summary.FileWriter(
                os.path.join(self._training_options.summary_directory, "cv"))