diff --git a/src/detector.c b/src/detector.c index baedab5f..021b3008 100644 --- a/src/detector.c +++ b/src/detector.c @@ -221,6 +221,10 @@ void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, i avg_loss = avg_loss*.9 + loss*.1; i = get_current_batch(net); + if (net.cudnn_half) { + if (i < net.burn_in) printf("\n Tensor Cores are disabled until the first %d iterations are reached.", net.burn_in); + else printf("\n Tensor Cores are used."); + } printf("\n %d: %f, %f avg loss, %f rate, %lf seconds, %d images\n", get_current_batch(net), loss, avg_loss, get_current_rate(net), (what_time_is_it_now()-time), i*imgs); #ifdef OPENCV diff --git a/src/network.h b/src/network.h index f630fb96..65029947 100644 --- a/src/network.h +++ b/src/network.h @@ -145,7 +145,7 @@ YOLODLL_API void reset_rnn(network *net); YOLODLL_API network *load_network_custom(char *cfg, char *weights, int clear, int batch); YOLODLL_API network *load_network(char *cfg, char *weights, int clear); YOLODLL_API float *network_predict_image(network *net, image im); -YOLODLL_API void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, int ngpus, int clear, int dont_show); +YOLODLL_API void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, int ngpus, int clear, int dont_show, int calc_map); YOLODLL_API int network_width(network *net); YOLODLL_API int network_height(network *net);