From 7a20e5dfa0d7a9362e242108a627b2cc7638d6e0 Mon Sep 17 00:00:00 2001
From: DEYIWANG <49895510+deyiwang89@users.noreply.github.com>
Date: Wed, 20 Apr 2022 09:25:01 +0800
Subject: [PATCH] Update summary.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

修改yolobody调用方法
---
 summary.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/summary.py b/summary.py
index 7d3441f..204c849 100644
--- a/summary.py
+++ b/summary.py
@@ -1,13 +1,13 @@
-#--------------------------------------------#
-#   该部分代码用于看网络结构
-#--------------------------------------------#
 import torch
 from torchsummary import summary
 
-from nets.yolo import YoloBody
+from nets.yolo4_tiny import YoloBody
+
+import numpy as np
+
 
 if __name__ == "__main__":
     # 需要使用device来指定网络在GPU还是CPU运行
     device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    m       = YoloBody([[3, 4, 5], [1, 2, 3]], 80).to(device)
-    summary(m, input_size=(3, 416, 416))
+    m       = YoloBody(2).to(device)
+    summary(m, input_data=(3, 416, 416))