Python SummaryWriter.add_embedding方法代码示例

# 需要导入模块: from tensorboardX import SummaryWriter [as 别名]
# 或者: from tensorboardX.SummaryWriter import add_embedding [as 别名]

#.........这里部分代码省略.........
    attribute_names = [train_dataset.get_attribute(al).name
                       for al in all_attribute_labels]

    # Initialize the model
    model = _make_cuda(SemanticAutoencoder(
        input_image_size, patch_size, num_attributes))

    # Initialize the loss function and optimizer
    epoch_loss = None
    criterion = _make_cuda(CustomLoss2(lambda_val=loss_lambda))
    optimizer = optim.Adam(ifilter(lambda p: p.requires_grad,
                                   model.parameters()))

    # Initiate training
    pbar, steps = tqdm(range(1, num_epochs + 1)), 0
    for epoch in pbar:
        epoch_loss = 0.

        model.train()  # Setting the model in training mode for training
        for image, label, attribute_labels, padding_idx in train_dataloader:
            steps += 1  # Incrementing the global step
            model.zero_grad()  # Clearing the gradients for each mini-batch

            # Create the input variable and get the output from the model
            x = _make_cuda(torch.autograd.Variable(image))
            z, z_patches, reconstructed_x = model(x)

            # Get the associated prototypes for each image in the batch
            prototype_labels = _make_cuda(attribute_labels)
            positive_prototypes = model.prototypes(prototype_labels)

            # Get the *non-associated* prototypes for each image in the batch
            negative_prototypes = list()
            for img_al in attribute_labels:
                negative_al = _make_cuda(torch.LongTensor(list(filter(
                    lambda al: al not in img_al,
                    all_attribute_labels))))
                negative_prototypes.append(model.prototypes(negative_al))

            # Compute the loss
            loss = criterion(reconstructed_x, z_patches,
                             positive_prototypes, padding_idx, x,
                             negative_prototypes=negative_prototypes)

            # Do backprop and update the weights
            loss.backward()
            optimizer.step()

            # Update the epoch loss and add the step loss to tensorboard
            epoch_loss += loss.item()
            writer.add_scalar('loss/step_loss', loss, steps)

        # Add the epoch loss to tensorboard and update the progressbar
        writer.add_scalar('loss/epoch_loss', epoch_loss, steps)
        pbar.set_postfix(epoch_loss=epoch_loss)

        model.eval()  # Setting the model in evaluation mode for testing
        if (epoch % 5 == 0) or (epoch == num_epochs):
            # Compute the nearest patch for each prototype
            nearest_patches_for_prototypes = \
                model.get_nearest_patches_for_prototypes(
                    train_dataset_with_non_random_transformation)

            # Update each prototype to be equal to the nearest patch
            model.reproject_prototypes(nearest_patches_for_prototypes)

            if (epoch % 1000 == 0) or (epoch == num_epochs):
                # Save the prototype visualization
                save_prototype_patch_visualization(
                    model, train_dataset_with_non_random_transformation,
                    nearest_patches_for_prototypes, PROTOTYPES_DIR)

                # Save the reconstructed images for the test dataset
                # for every 1000 epochs
                for i_, (image, image_label, attribute_labels, _) \
                        in enumerate(test_dataset):
                    x = image.view((1,) + image.size())
                    x = _make_cuda(torch.autograd.Variable(x))
                    z, z_patches, reconstructed_x = model(x)

                    reconstructed_image = \
                        get_image_from_tensor(reconstructed_x)
                    reconstructed_image.save(
                        os.path.join(IMAGES_DIR, '%d-%d.png' % (epoch, i_)))

                # Save the intermediate model
                model.save_weights(os.path.join(RUN_DIR, MODEL_FILE_NAME))

        # Add the prototype embeddings to tensorboard at the end
        if epoch == num_epochs:
            writer.add_embedding(
                model.prototypes.weight[1:],
                metadata=attribute_names,
                global_step=steps)

    # Save the final model and commit the tensorboard logs
    model.save_weights(os.path.join(RUN_DIR, MODEL_FILE_NAME))
    writer.close()

    return epoch_loss

# 需要导入模块: from tensorboardX import SummaryWriter [as 别名]
# 或者: from tensorboardX.SummaryWriter import add_embedding [as 别名]

#.........这里部分代码省略.........
        :param comment: Descriptive comment to append to path
        :type comment: str
        :param num_images: The number of images to write
        :type num_images: int
        :param avg_pool_size: Size of the average pool to perform on the image. This is recommended to reduce the
        overall image sizes and improve latency
        :type avg_pool_size: int
        :param avg_data_channels: If True, the image data will be averaged in the channel dimension
        :type avg_data_channels: bool
        :param write_data: If True, the raw data will be written as an embedding
        :type write_data: bool
        :param write_features: If True, the image features will be written as an embedding
        :type write_features: bool
        :param features_key: The key in state to use for the embedding. Typically model output but can be used to show
        features from any layer of the model.
        :type features_key: str
        """
        self.log_dir = log_dir
        self.comment = comment
        self.num_images = num_images
        self.avg_pool_size = avg_pool_size
        self.avg_data_channels = avg_data_channels
        self.write_data = write_data
        self.write_features = write_features
        self.features_key = features_key

        self._writer = None

        self.done = False

    def on_start(self, state):
        log_dir = os.path.join(self.log_dir, state[torchbearer.MODEL].__class__.__name__ + '_' + self.comment)
        self._writer = SummaryWriter(log_dir=log_dir)

    def on_step_validation(self, state):
        if not self.done:
            x = state[torchbearer.X].data.clone()

            if len(x.size()) == 3:
                x = x.unsqueeze(1)

            x = F.avg_pool2d(x, self.avg_pool_size).data

            data = None

            if state[torchbearer.EPOCH] == 0 and self.write_data:
                if self.avg_data_channels:
                    data = torch.mean(x, 1)
                else:
                    data = x

                data = data.view(data.size(0), -1)

            feature = None

            if self.write_features:
                feature = state[self.features_key].data.clone()
                feature = feature.view(feature.size(0), -1)

            label = state[torchbearer.Y_TRUE].data.clone()

            if state[torchbearer.BATCH] == 0:
                remaining = self.num_images if self.num_images < label.size(0) else label.size(0)

                self._images = x[:remaining].to('cpu')
                self._labels = label[:remaining].to('cpu')

                if data is not None:
                    self._data = data[:remaining].to('cpu')

                if feature is not None:
                    self._features = feature[:remaining].to('cpu')
            else:
                remaining = self.num_images - self._labels.size(0)

                if remaining > label.size(0):
                    remaining = label.size(0)

                self._images = torch.cat((self._images, x[:remaining].to('cpu')), dim=0)
                self._labels = torch.cat((self._labels, label[:remaining].to('cpu')), dim=0)

                if data is not None:
                    self._data = torch.cat((self._data, data[:remaining].to('cpu')), dim=0)

                if feature is not None:
                    self._features = torch.cat((self._features, feature[:remaining].to('cpu')), dim=0)

            if self._labels.size(0) >= self.num_images:
                if state[torchbearer.EPOCH] == 0 and self.write_data:
                    self._writer.add_embedding(self._data, metadata=self._labels, label_img=self._images, tag='data', global_step=-1)
                if self.write_features:
                    self._writer.add_embedding(self._features, metadata=self._labels, label_img=self._images, tag='features', global_step=state[torchbearer.EPOCH])
                self.done = True

    def on_end_epoch(self, state):
        if self.write_features:
            self.done = False

    def on_end(self, state):
        self._writer.close()