Преглед на файлове

`yolo.py --profile` updates (#7170)

modifyDataloader
Glenn Jocher GitHub преди 2 години
родител
ревизия
b2194b9015
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
променени са 1 файла, в които са добавени 2 реда и са изтрити 2 реда
  1. +2
    -2
      models/yolo.py

+ 2
- 2
models/yolo.py Целия файл

@@ -310,11 +310,11 @@ if __name__ == '__main__':

# Create model
model = Model(opt.cfg).to(device)
model.train()

# Profile
if opt.profile:
img = torch.rand(16 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
model.eval().fuse()
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
y = model(img, profile=True)

# Test all models

Loading…
Отказ
Запис