dcnn.cpp

   1 #define DEBUG
   2 #include <assert.h>
   3 #include <ctype.h>
   4 #include <math.h>
   5 #include <stdarg.h>
   6 #include <stdio.h>
   7 #include <stdlib.h>
   8 #include <string.h>
   9 #include <unistd.h>
  10 #include <sys/stat.h>
  11
  12 #define CPU_ONLY 1
  13 #include <caffe/caffe.hpp>
  14 using namespace caffe;
  15
  16 extern "C" {
  17 #include "debug.h"
  18 #include "board.h"
  19 #include "dcnn.h"
  20
  21
  22 static shared_ptr<Net<float> > net;
  23
  24 bool
  25 using_dcnn(struct board *b)
  26 {
  27         return (real_board_size(b) == 19 && net);
  28 }
  29
  30 /* Make caffe quiet */
  31 void
  32 dcnn_quiet_caffe(int argc, char *argv[])
  33 {
  34         if (DEBUGL(7) || getenv("GLOG_minloglevel"))
  35                 return;
  36
  37         setenv("GLOG_minloglevel", "2", 1);
  38         execvp(argv[0], argv);   /* Sucks that we have to do this */
  39 }
  40
  41 void
  42 dcnn_init()
  43 {
  44         if (net)
  45                 return;
  46
  47         struct stat s;
  48         const char *model_file =   "movepredict.prototxt";
  49         const char *trained_file = "movepredict.caffemodel";
  50         if (stat(model_file, &s) != 0  ||  stat(trained_file, &s) != 0) {
  51                 if (DEBUGL(1))
  52                         fprintf(stderr, "No dcnn files found, will not use dcnn code.\n");
  53                 return;
  54         }
  55
  56         Caffe::set_mode(Caffe::CPU);
  57
  58         /* Load the network. */
  59         net.reset(new Net<float>(model_file, TEST));
  60         net->CopyTrainedLayersFrom(trained_file);
  61
  62         if (DEBUGL(1))
  63                 fprintf(stderr, "Initialized dcnn.\n");
  64 }
  65
  66 void
  67 dcnn_get_moves(struct board *b, enum stone color, float result[])
  68 {
  69         assert(real_board_size(b) == 19);
  70         int size = 19;
  71         float *data = new float[2 * size * size];
  72
  73         if (color == S_BLACK)
  74                 for (int j = 0; j < size; j++)
  75                         for (int k = 0; k < size; k++) {
  76                                 coord_t c = coord_xy(b, j + 1, k + 1);
  77                                 int offsetb = j * size + k;
  78                                 int offsetw = size * size + size * j + k;
  79                                 data[offsetb] = (board_at(b, c) == S_BLACK ? 1.0 : 0.0);
  80                                 data[offsetw] = (board_at(b, c) == S_WHITE ? 1.0 : 0.0);
  81                         }
  82
  83         if (color == S_WHITE) {
  84                 for (int j = 0; j < size; j++)
  85                         for (int k = 0; k < size; k++) {
  86                                 coord_t c = coord_xy(b, j + 1, k + 1);
  87                                 int offsetb = j * size + k;
  88                                 int offsetw = size * size + size * j + k;
  89                                 data[offsetb] = (board_at(b, c) == S_WHITE ? 1.0 : 0.0);
  90                                 data[offsetw] = (board_at(b, c) == S_BLACK ? 1.0 : 0.0);
  91                         }
  92         }
  93
  94         Blob<float> *blob = new Blob<float>(1,2,size,size);
  95         blob->set_cpu_data(data);
  96         vector<Blob<float>*> bottom;
  97         bottom.push_back(blob);
  98         assert(net);
  99         const vector<Blob<float>*>& rr = net->Forward(bottom);
 100
 101         for (int i = 0; i < size * size; i++) {
 102                 result[i] = rr[0]->cpu_data()[i];
 103                 if (result[i] < 0.00001)
 104                         result[i] = 0.00001;
 105         }
 106         delete[] data;
 107         delete blob;
 108 }
 109
 110
 111 } /* extern "C" */
 112