소스 검색

`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…
취소
저장