From 4d93b07469179d98d3ce67fb5c8ee0e1173ca585 Mon Sep 17 00:00:00 2001 From: Anner Date: Tue, 2 Oct 2018 15:10:54 +0900 Subject: [PATCH] Simplify Darknet() class: forward pass yolo case You mentioned in your blog we can't concatenate the first yolo layer detections to an empty tensor, but I think we actually can. This simplifies your code a bit. Also, not 100% sure so didnt delete this, but I think you can directly send x to predict_transform(), why would you send x.data instead? And why is there the `if type(x) == int: continue`? predict_transform() will never return an int right? last, why: `outputs[i] = outputs[i-1]`? This might be useful for future different use cases, but in the yolo architecture you no routing layer or shortcut layer will ever need to access the yolo predictions, so no need to keep them around in outputs dict. --- darknet.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/darknet.py b/darknet.py index facb5d75..bbb7b910 100644 --- a/darknet.py +++ b/darknet.py @@ -305,12 +305,10 @@ def get_module_list(self): def forward(self, x, CUDA): - detections = [] modules = self.blocks[1:] outputs = {} #We cache the outputs for the route layer + detections = torch.empty(0) # concatenate all bbox predictions to this initially empty tensor - - write = 0 for i in range(len(modules)): module_type = (modules[i]["type"]) @@ -364,17 +362,7 @@ def forward(self, x, CUDA): if type(x) == int: continue - - if not write: - detections = x - write = 1 - - else: - detections = torch.cat((detections, x), 1) - - outputs[i] = outputs[i-1] - - + detections = torch.cat((detections, x), dim=1) try: return detections