lichen/detecttracking
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cpu to device select

Browse Source
This commit is contained in:
王庆刚
2025-01-13 18:11:56 +08:00
parent a50f777839
commit 744fb7b7b2
4 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
__pycache__/track_reid.cpython-39.pyc
View File

Binary file not shown.

4
contrast/feat_extract/config.py
View File

@ -76,8 +76,8 @@ class Config:
lr_decay = 0.95 # 0.98
weight_decay = 5e-4
loss = 'cross_entropy' # ['focal_loss', 'cross_entropy']
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
pin_memory = True # if memory is large, set it True to speed up a bit
num_workers = 4 # dataloader

BIN
models/__pycache__/experimental.cpython-39.pyc
View File

Binary file not shown.

2
models/experimental.py
View File

@ -76,7 +76,7 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
ckpt = torch.load(attempt_download(w), map_location='cpu') # load
ckpt = torch.load(attempt_download(w), map_location=device) # load
ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model
# Model compatibility updates