This commit is contained in:
Glenn Jocher
2019-02-12 16:58:07 +01:00
parent e4e64a9ff6
commit cc5e9a5a85
5 changed files with 5 additions and 12 deletions
Executable → Regular
+1
View File
@@ -1,5 +1,6 @@
import argparse import argparse
import time import time
from sys import platform
from models import * from models import *
from utils.datasets import * from utils.datasets import *
+1 -1
View File
@@ -243,7 +243,7 @@ class Darknet(nn.Module):
def __init__(self, cfg_path, img_size=416): def __init__(self, cfg_path, img_size=416):
super(Darknet, self).__init__() 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]['cfg'] = cfg_path
self.module_defs[0]['height'] = img_size self.module_defs[0]['height'] = img_size
self.hyperparams, self.module_list = create_modules(self.module_defs) self.hyperparams, self.module_list = create_modules(self.module_defs)
-1
View File
@@ -2,7 +2,6 @@ import glob
import math import math
import os import os
import random import random
from sys import platform
import cv2 import cv2
import numpy as np import numpy as np
+1 -1
View File
@@ -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""" """Parses the yolo-v3 layer configuration file and returns module definitions"""
file = open(path, 'r') file = open(path, 'r')
lines = file.read().split('\n') lines = file.read().split('\n')
+2 -9
View File
@@ -1,16 +1,9 @@
import torch import torch
def check_cuda():
return torch.cuda.is_available()
CUDA_AVAILABLE = check_cuda()
def init_seeds(seed=0): def init_seeds(seed=0):
torch.manual_seed(seed) torch.manual_seed(seed)
if CUDA_AVAILABLE: if torch.cuda.is_available():
torch.cuda.manual_seed(seed) torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available # 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: if force_cpu:
device = torch.device('cpu') device = torch.device('cpu')
else: 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') print('Using ' + str(device) + '\n')
return device return device