Method probe_data

modules/ace_plus_solver.py:130–181

(self)

        return log_data, log_label
    @property
    def probe_data(self):
        if not we.debug and self.mode == 'train':
            batch_data = transfer_data_to_cuda(self.current_batch_data[self.mode])
            self.eval_mode()
            with torch.autocast(device_type='cuda',
                                enabled=self.use_amp,
                                dtype=self.dtype):
                batch_data['log_num'] = self.log_train_num
                batch_data.update(self.sample_args.get_lowercase_dict())
                results = self.run_step_eval(batch_data)
            self.train_mode()
            log_data, log_label = self.save_results(results)
            self.register_probe({
                'train_image':
                    ProbeData(log_data,
                              is_image=True,
                              build_html=True,
                              build_label=log_label)
            })
            self.register_probe({'train_label': log_label})
            if self.probe_prompt:
                self.eval_mode()
                all_results = []
                for prompt in self.probe_prompt:
                    with torch.autocast(device_type='cuda',
                                        enabled=self.use_amp,
                                        dtype=self.dtype):
                        batch_data = {
                            "prompt": [[prompt]],
                            "image": [torch.zeros(3, self.probe_hw[0], self.probe_hw[1])],
                            "image_mask": [torch.ones(1, self.probe_hw[0], self.probe_hw[1])],
                            "src_image_list": [[]],
                            "modify_image_list": [[]],
                            "src_mask_list": [[]],
                            "edit_id": [[]],
                            "height": self.probe_hw[0],
                            "width": self.probe_hw[1]
                        }
                        batch_data.update(self.sample_args.get_lowercase_dict())
                        results = self.run_step_eval(batch_data)
                        all_results.extend(results)
                self.train_mode()
                log_data, log_label = self.save_results(all_results)
                self.register_probe({
                    'probe_image':
                        ProbeData(log_data,
                                  is_image=True,
                                  build_html=True,
                                  build_label=log_label)
                })

        return super(LatentDiffusionSolver, self).probe_data