How to reduce the time to test a single image?

[x12901]

How to reduce the time to test a single image
I want to test a single image, but I have to load the model once every time, which is a waste of time. I load the model in on_test_start() and predict the images in test_step(). Is it possible to load the image directly into test_step() without going to on_test_start() ？

class STPM(pl.LightningModule):
    def __init__(self, hparams):
        ...
        
    def forward(self, x_t):
        ...

    def test_dataloader(self):
        test_datasets = MVTecDataset(root=os.path.join(args.dataset_path, args.category),
                                     transform=self.data_transforms, phase='test')
        test_loader = DataLoader(test_datasets, batch_size=1, shuffle=False,
                                 num_workers=0) 
        return test_loader


    def on_test_start(self):
        self.init_results_list()
        self.embedding_dir_path, self.sample_path, self.source_code_save_path = prep_dirs(self.logger.log_dir)
        self.embedding_coreset = pickle.load(open(os.path.join(self.embedding_dir_path, 'embedding.pickle'), 'rb'))
        embeded = torch.tensor(self.embedding_coreset)
        train_jit = TrainFeature(embeded)
        traced_model = torch.jit.script(train_jit)
        torch.jit.save(traced_model, "patchcore_features.pt")
        dummy_input = torch.randn(args.batch_size, 3, args.input_size, args.input_size, requires_grad=False).cuda()
        traced_script_module = torch.jit.trace(self.model, dummy_input)
        traced_script_module.save("patchcore_model.pt")

    def test_step(self, batch, batch_idx):  # Nearest Neighbour Search
        ...

    def test_epoch_end(self, outputs):
        self.analyze_data()
        best_th, threshold = self.Find_Optimal_Cutoff(self.y_true, self.all_scores_mean_norm)
        print(f'\nbest threshold={best_th}')
        ng_index = np.where(np.array(self.y_true) == 1)
        if len(ng_index[0]) == 0:
            ng_index = len(self.y_true)
        else:
            ng_index = ng_index[0][0]
        fig = plt.figure()
        sns.histplot(self.all_scores_mean_norm[:ng_index], kde=True, color="blue", label="normal")
        sns.histplot(self.all_scores_mean_norm[ng_index:], kde=True, color="red", label="abnormal")
        fig.legend(labels=['normal', 'abnormal'])
        plt.xlabel("Anomaly score")
        plt.ylabel("Count")
        plt.savefig('Anomaly_score_histplot.jpg')


def get_args():
    parser = argparse.ArgumentParser(description='ANOMALYDETECTION')
    parser.add_argument('--phase', choices=['train', 'test'], default='test')
    parser.add_argument('--dataset_path', default=r'C:\Users\Administrator\Desktop\dataset')  
    parser.add_argument('--category', default='a') 
    parser.add_argument('--num_epochs', default=1)
    parser.add_argument('--batch_size', default=1)
    parser.add_argument('--load_size', default=256)  
    parser.add_argument('--input_size', default=224)  
    parser.add_argument('--coreset_sampling_ratio', default=0.001)
    parser.add_argument('--project_root_path',
                        default=r'D:\code')  
    parser.add_argument('--save_src_code', default=True)
    parser.add_argument('--save_anomaly_map', default=True)
    parser.add_argument('--n_neighbors', type=int, default=9)
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args = get_args()
    trainer = pl.Trainer.from_argparse_args(args, default_root_dir=os.path.join(args.project_root_path, args.category),
                                            max_epochs=args.num_epochs,
                                            gpus=1)
    model = STPM(hparams=args)
    if args.phase == 'train':
        trainer.fit(model)
        trainer.test(model)
    elif args.phase == 'test':
        trainer.test(model)

[tchaton]

Dear @x12901,

I am not fully sure to grasp your problem. Would you mind explaining more in details where the time is lost and what you are trying to do. Do you have only 1 image in your test dataset ?

Best,
T.C