未验证 提交 7a621d24 编写于 作者: B Bubbliiiing 提交者: GitHub

Add files via upload

上级 c530cfe1
import torch
from torchsummary import summary
from nets.CSPdarknet import darknet53
from nets.yolo4 import YoloBody
#--------------------------------------------#
# 该部分代码只用于看网络结构,并非测试代码
# map测试请看get_dr_txt.py、get_gt_txt.py
# 和get_map.py
#--------------------------------------------#
if __name__ == "__main__":
# 需要使用device来指定网络在GPU还是CPU运行
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = YoloBody(3,20).to(device)
summary(model, input_size=(3, 416, 416))
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册