Skip to content

GitLab

Commit dcb5e4e4 authored by wangkx1's avatar wangkx1
Browse files

Update torch_verify.py

parent 32eb2157
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torch_verify.py
View file @ dcb5e4e4
...@@ -7,7 +7,7 @@ from PIL import Image ...@@ -7,7 +7,7 @@ from PIL import Image
# 加载预训练的ResNet50模型，并指定不使用GPU # 加载预训练的ResNet50模型，并指定不使用GPU
model = models.resnet50(pretrained=True) model = models.resnet50(pretrained=True)
model.eval() # 将模型设置为评估模式 model.eval() # 将模型设置为评估模式
device = torch.device("cpu") # 指定设备为CPU device = torch.device("cuda:0") # 指定设备为CPU
model.to(device) # 将模型转移到CPU上 model.to(device) # 将模型转移到CPU上
# 图片预处理 # 图片预处理
...@@ -45,4 +45,4 @@ with open(synset_words_path, 'r') as f: ...@@ -45,4 +45,4 @@ with open(synset_words_path, 'r') as f:
# 每行前去除空格和制表符，并根据索引获取类别名称 # 每行前去除空格和制表符，并根据索引获取类别名称
class_names = [line.strip() for line in synset_words] class_names = [line.strip() for line in synset_words]
predicted_class_name = class_names[predicted_class_idx+1] # 注意索引可能从1开始，所以+1 predicted_class_name = class_names[predicted_class_idx+1] # 注意索引可能从1开始，所以+1
print(f"Predicted class name: {predicted_class_name}") # 彭布罗克，彭布罗克威尔士柯基犬 print(f"Predicted class name: {predicted_class_name}") # 彭布罗克，彭布罗克威尔士柯基犬
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment