goodfellowliu commited on
Commit
2368603
·
1 Parent(s): a557b7d

fix refrence bug

Browse files

In torch==1.5, the import of the API has changed. Although it does not interrupt the operation of the program, it seems to me to be an implicit error and may throw an exception in later versions.

Files changed (2) hide show
  1. detect.py +3 -1
  2. train.py +1 -0
detect.py CHANGED
@@ -1,5 +1,7 @@
1
  import argparse
2
 
 
 
3
  from utils.datasets import *
4
  from utils.utils import *
5
 
@@ -36,7 +38,7 @@ def detect(save_img=False):
36
  vid_path, vid_writer = None, None
37
  if webcam:
38
  view_img = True
39
- torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
40
  dataset = LoadStreams(source, img_size=imgsz)
41
  else:
42
  save_img = True
 
1
  import argparse
2
 
3
+ import torch.backends.cudnn as cudnn
4
+
5
  from utils.datasets import *
6
  from utils.utils import *
7
 
 
38
  vid_path, vid_writer = None, None
39
  if webcam:
40
  view_img = True
41
+ cudnn.benchmark = True # set True to speed up constant image size inference
42
  dataset = LoadStreams(source, img_size=imgsz)
43
  else:
44
  save_img = True
train.py CHANGED
@@ -4,6 +4,7 @@ import torch.distributed as dist
4
  import torch.nn.functional as F
5
  import torch.optim as optim
6
  import torch.optim.lr_scheduler as lr_scheduler
 
7
  from torch.utils.tensorboard import SummaryWriter
8
 
9
  import test # import test.py to get mAP after each epoch
 
4
  import torch.nn.functional as F
5
  import torch.optim as optim
6
  import torch.optim.lr_scheduler as lr_scheduler
7
+ import torch.utils.data
8
  from torch.utils.tensorboard import SummaryWriter
9
 
10
  import test # import test.py to get mAP after each epoch