Bläddra i källkod

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

modifyDataloader
Glenn Jocher GitHub 2 år sedan
förälder
incheckning
b2194b9015
Ingen känd nyckel hittad för denna signaturen i databasen GPG-nyckel ID: 4AEE18F83AFDEB23
1 ändrade filer med 2 tillägg och 2 borttagningar
  1. +2
    -2
      models/yolo.py

+ 2
- 2
models/yolo.py Visa fil

@@ -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

Laddar…
Avbryt
Spara