Mathis/CUDANet
1
0
Fork 0
mirror of https://github.com/lordmathis/CUDANet.git synced 2025-11-06 01:34:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update inception v3 readme

Browse Source
This commit is contained in:
Matúš Námešný
2024-09-04 21:32:05 +02:00
parent c8bc6f7a39
commit 7b8c4bd811
3 changed files with 56 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tools/utils.py
View File

@@ -35,8 +35,6 @@ def export_model_weights(model: torch.nn.Module, filename):
tensor_data += tensor_bytes
# print(model.named_buffers)
# Add buffers (for running_mean and running_var)
for name, buf in model.named_buffers():
if "running_mean" not in name and "running_var" not in name:
@@ -76,9 +74,7 @@ def predict(model, image_path, resize=299, crop=299, preprocess=None):
)
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(
0
) # create a mini-batch as expected by the model
input_batch = input_tensor.unsqueeze(0)
# move the input and model to GPU for speed if available
if torch.cuda.is_available():
@@ -86,5 +82,5 @@ def predict(model, image_path, resize=299, crop=299, preprocess=None):
model.to("cuda")
with torch.no_grad():
output = model(input_batch)
return torch.argmax(output)
output = model(input_batch)
return torch.argmax(output).item()