Просмотр исходного кода

`Ensemble()` visualize fix (#3973)

* fix visualize error

* Revert "fix visualize error"

* add visualise profile
modifyDataloader
KEN GitHub 3 лет назад
Родитель
Сommit
647223a7a8
Не найден GPG ключ соответствующий данной подписи Идентификатор GPG ключа: 4AEE18F83AFDEB23
1 измененных файлов: 2 добавлений и 2 удалений
  1. +2
    -2
      models/experimental.py

+ 2
- 2
models/experimental.py Просмотреть файл

@@ -100,10 +100,10 @@ class Ensemble(nn.ModuleList):
def __init__(self):
super(Ensemble, self).__init__()

def forward(self, x, augment=False):
def forward(self, x, augment=False, profile=False, visualize=False):
y = []
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).mean(0) # mean ensemble
y = torch.cat(y, 1) # nms ensemble

Загрузка…
Отмена
Сохранить