updates
This commit is contained in:
parent
e4e64a9ff6
commit
cc5e9a5a85
|
@ -1,5 +1,6 @@
|
|||
import argparse
|
||||
import time
|
||||
from sys import platform
|
||||
|
||||
from models import *
|
||||
from utils.datasets import *
|
||||
|
|
|
@ -243,7 +243,7 @@ class Darknet(nn.Module):
|
|||
def __init__(self, cfg_path, img_size=416):
|
||||
super(Darknet, self).__init__()
|
||||
|
||||
self.module_defs = parse_model_config(cfg_path)
|
||||
self.module_defs = parse_model_cfg(cfg_path)
|
||||
self.module_defs[0]['cfg'] = cfg_path
|
||||
self.module_defs[0]['height'] = img_size
|
||||
self.hyperparams, self.module_list = create_modules(self.module_defs)
|
||||
|
|
|
@ -2,7 +2,6 @@ import glob
|
|||
import math
|
||||
import os
|
||||
import random
|
||||
from sys import platform
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
|
|
@ -1,4 +1,4 @@
|
|||
def parse_model_config(path):
|
||||
def parse_model_cfg(path):
|
||||
"""Parses the yolo-v3 layer configuration file and returns module definitions"""
|
||||
file = open(path, 'r')
|
||||
lines = file.read().split('\n')
|
||||
|
|
|
@ -1,16 +1,9 @@
|
|||
import torch
|
||||
|
||||
|
||||
def check_cuda():
|
||||
return torch.cuda.is_available()
|
||||
|
||||
|
||||
CUDA_AVAILABLE = check_cuda()
|
||||
|
||||
|
||||
def init_seeds(seed=0):
|
||||
torch.manual_seed(seed)
|
||||
if CUDA_AVAILABLE:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
|
||||
|
@ -20,6 +13,6 @@ def select_device(force_cpu=False):
|
|||
if force_cpu:
|
||||
device = torch.device('cpu')
|
||||
else:
|
||||
device = torch.device('cuda:0' if CUDA_AVAILABLE else 'cpu')
|
||||
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||||
print('Using ' + str(device) + '\n')
|
||||
return device
|
||||
|
|
Loading…
Reference in New Issue