|
|
@@ -18,7 +18,7 @@ import torch.distributed as dist |
|
|
|
import torch.nn as nn |
|
|
|
import torch.nn.functional as F |
|
|
|
|
|
|
|
from utils.general import LOGGER, file_update_date, git_describe |
|
|
|
from utils.general import LOGGER, file_date, git_describe |
|
|
|
|
|
|
|
try: |
|
|
|
import thop # for FLOPs computation |
|
|
@@ -51,7 +51,7 @@ def device_count(): |
|
|
|
|
|
|
|
def select_device(device='', batch_size=0, newline=True): |
|
|
|
# device = 'cpu' or '0' or '0,1,2,3' |
|
|
|
s = f'YOLOv5 🚀 {git_describe() or file_update_date()} torch {torch.__version__} ' # string |
|
|
|
s = f'YOLOv5 🚀 {git_describe() or file_date()} Python-{platform.python_version()} torch-{torch.__version__} ' |
|
|
|
device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0' |
|
|
|
cpu = device == 'cpu' |
|
|
|
if cpu: |