Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Update models.py

This commit is contained in:
Glenn Jocher 2019-03-05 18:43:51 +01:00 committed by GitHub
parent 6fb14fc903
commit 473eb8d0c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
models.py
View File

@ -265,6 +265,11 @@ class Darknet(nn.Module):
return sum(output) if is_training else torch.cat(output, 1) return sum(output) if is_training else torch.cat(output, 1)
def get_yolo_layers(model):
a = [module_def['type'] == 'yolo' for module_def in model.module_defs]
return [i for i, x in enumerate(a) if x] # [82, 94, 106] for yolov3
def create_grids(self, img_size, nG): def create_grids(self, img_size, nG):
self.stride = img_size / nG self.stride = img_size / nG