From 266bcc065128cc05e80dde5e6bb1c00acb4c0ba2 Mon Sep 17 00:00:00 2001
From: ChengZi <chen.zhang@zilliz.com>
Date: Fri, 18 Mar 2022 17:44:51 +0800
Subject: [PATCH] add converting to cpu

---
 default_config.yaml       | 38 --------------------------------------
 resnet_image_embedding.py |  3 +++
 2 files changed, 3 insertions(+), 38 deletions(-)
 delete mode 100644 default_config.yaml

diff --git a/default_config.yaml b/default_config.yaml
deleted file mode 100644
index 8c1d885..0000000
--- a/default_config.yaml
+++ /dev/null
@@ -1,38 +0,0 @@
-train:
-    output_dir: ./output_dir
-    overwrite_output_dir: true
-    eval_strategy: epoch
-    eval_steps:
-    batch_size: 8
-    val_batch_size: -1
-    seed: 42
-    epoch_num: 2
-    dataloader_pin_memory: true
-    dataloader_drop_last: true
-    dataloader_num_workers: 0
-    load_best_model_at_end: false
-    freeze_bn: false
-learning:
-    lr: 5e-05
-    loss: CrossEntropyLoss
-    optimizer: Adam
-    lr_scheduler_type: linear
-    warmup_ratio: 0.0
-    warmup_steps: 0
-callback:
-    early_stopping:
-        monitor: eval_epoch_metric
-        patience: 4
-        mode: max
-    model_checkpoint:
-        every_n_epoch: 1
-    tensorboard:
-        log_dir:
-        comment: ''
-logging:
-    print_steps:
-metrics:
-    metric: Accuracy
-device:
-    device_str:
-    sync_bn: false
diff --git a/resnet_image_embedding.py b/resnet_image_embedding.py
index 985e57c..e59f543 100644
--- a/resnet_image_embedding.py
+++ b/resnet_image_embedding.py
@@ -34,6 +34,7 @@ class ResnetImageEmbedding(NNOperator):
     def __init__(self, model_name: str, num_classes: int = 1000, framework: str = 'pytorch') -> None:
         super().__init__(framework=framework)
         self.model = timm.create_model(model_name, pretrained=True)
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         pretrained_dict = None
         if model_name == 'resnet101':
             pretrained_dict = torch.hub.load_state_dict_from_url(
@@ -53,12 +54,14 @@ class ResnetImageEmbedding(NNOperator):
                                         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
 
     def __call__(self, image: 'towhee.types.Image') -> NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]):
+        self.model.to(self.device)
         img_tensor = self.tfms(to_pil(image)).unsqueeze(0)
         self.model.eval()
         features = self.model.forward_features(img_tensor)
         if features.dim() == 4:  # if the shape of feature map is [N, C, H, W], where H > 1 and W > 1
             global_pool = nn.AdaptiveAvgPool2d(1)
             features = global_pool(features)
+        features = features.to('cpu')
         embedding = features.flatten().detach().numpy()
         Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)])
         return Outputs(embedding)