def run(self):
"""This method runs in the new process."""
global logger
setup_exceptions()
logger = log_utils.setup_logger('tile_worker')
if self.caffe_path is not None:
sys.path.append(self.caffe_path + '/python')
if self.device >= 0:
os.environ['CUDA_VISIBLE_DEVICES'] = str(self.device)
import caffe
if self.device >= 0:
caffe.set_mode_gpu()
else:
caffe.set_mode_cpu()
caffe.set_random_seed(0)
np.random.seed(0)
self.model = CaffeModel(*self.model_info)
self.model.img = np.zeros((3, 1, 1), dtype=np.float32)
while True:
try:
self.process_one_request()
except KeyboardInterrupt:
break
def init_model(resp_q, caffe_path, model, weights, mean):
"""Puts the list of layer shapes into resp_q. To be run in a separate process."""
global logger
setup_exceptions()
logger = log_utils.setup_logger('init_model')
if caffe_path:
sys.path.append(caffe_path + '/python')
import caffe
caffe.set_mode_cpu()
model = CaffeModel(model, weights, mean)
shapes = OrderedDict()
for layer in model.layers():
shapes[layer] = model.data[layer].shape
resp_q.put(shapes)
'--greedy',
help='Whether to use greedy stratagies.',
type=bool,
default=False)
parser.add_argument('-t',
'--timeout',
help='Timeout time.',
default=3,
type=int)
parser.add_argument('-r',
'--rety-time',
help='Retry time when connection failed.',
default=3,
type=int)
parser.add_argument('-n',
'--num-process',
help='Number of crawling process.',
default=1,
type=int)
parser.add_argument('-sf',
'--save-path',
help='Where to save results',
default=None)
if __name__ == '__main__':
args = parser.parse_args()
#logger = logging.getLogger()
#logger.setLevel(logging.INFO)
logger = setup_logger()
crawler(args)
return default
def setup_exceptions():
scheme = terminal_bg('LightBG', 'Linux', 'Neutral')
mode = 'Plain'
if 'DEBUG' in os.environ:
mode = 'Verbose'
try:
from IPython.core.ultratb import AutoFormattedTB
sys.excepthook = AutoFormattedTB(mode=mode, color_scheme=scheme)
except ImportError:
pass
logger = log_utils.setup_logger('style_transfer')
def set_thread_count(threads):
"""Sets the maximum number of MKL threads for this process."""
if MKL_THREADS is not None:
mkl.set_num_threads(max(1, threads))
try:
import mkl
MKL_THREADS = mkl.get_max_threads()
set_thread_count(1)
except ImportError:
pass