zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

`Ensemble()` visualize fix (#3973) 

* fix visualize error

* Revert "fix visualize error"

* add visualise profile
This commit is contained in:
KEN 2021-07-12 02:47:08 +09:00 committed by GitHub
parent a544d59f52
commit 647223a7a8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
models/experimental.py
View File

@ -100,10 +100,10 @@ class Ensemble(nn.ModuleList):
def __init__(self): def __init__(self):
super(Ensemble, self).__init__() super(Ensemble, self).__init__()
def forward(self, x, augment=False): def forward(self, x, augment=False, profile=False, visualize=False):
y = [] y = []
for module in self: for module in self:
y.append(module(x, augment)[0]) y.append(module(x, augment, profile, visualize)[0])
# y = torch.stack(y).max(0)[0] # max ensemble # y = torch.stack(y).max(0)[0] # max ensemble
# y = torch.stack(y).mean(0) # mean ensemble # y = torch.stack(y).mean(0) # mean ensemble
y = torch.cat(y, 1) # nms ensemble y = torch.cat(y, 1) # nms ensemble