1

data/generator.py

@@ -0,0 +1,40 @@
+import glob
+import torch
+from torch.nn.functional import one_hot
+
+
+class DataGenerator:
+    def __init__(self, data_directory, samples):
+        self.data_directory = data_directory
+        self.samples = samples
+        self.files = set(file_name for file_name, index in samples)
+        self.num_samples = None
+
+    def get_num_samples(self, batch_size=128, num_classes=19 * 19):
+        if self.num_samples is not None:
+            return self.num_samples
+        else:
+            self.num_samples = 0
+            for X, y in self._generate(batch_size=batch_size, num_classes=num_classes):
+                self.num_samples += X.shape[0]
+            return self.num_samples
+
+    def _generate(self, batch_size, num_classes):
+        for zip_file_name in self.files:
+            file_name = zip_file_name.replace('.tar.gz', '') + 'train'
+            base = self.data_directory + '/' + file_name + '_features_*.npy'
+            for feature_file in glob.glob(base):
+                label_file = feature_file.replace('features', 'labels')
+                x = torch.from_numpy(np.load(feature_file)).float()
+                y = torch.from_numpy(np.load(label_file)).long()
+                y = one_hot(y, num_classes=num_classes)
+                while x.shape[0] >= batch_size:
+                    x_batch, x = x[:batch_size], x[batch_size:]
+                    y_batch, y = y[:batch_size], y[batch_size:]
+                    yield x_batch, y_batch
+
+    def generate(self, batch_size=128, num_classes=19 * 19):
+        while True:
+            for item in self._generate(batch_size, num_classes):
+                yield item
+

data/index_processor.py

@@ -0,0 +1,107 @@
+# This Source Code Form is subject to the terms of the Mozilla Public License,
+# v. 2.0. If a copy of the MPL was not distributed with this file, You can
+# obtain one at http://mozilla.org/MPL/2.0/.
+from __future__ import print_function
+from __future__ import absolute_import
+import os
+import sys
+import multiprocessing
+import six
+if sys.version_info[0] == 3:
+    from urllib.request import urlopen, urlretrieve
+else:
+    from urllib import urlopen, urlretrieve
+
+
+def worker(url_and_target):  # Parallelize data download via multiprocessing
+    try:
+        (url, target_path) = url_and_target
+        print('>>> Downloading ' + target_path)
+        urlretrieve(url, target_path)
+    except (KeyboardInterrupt, SystemExit):
+        print('>>> Exiting child process')
+
+
+class KGSIndex:
+
+    def __init__(self,
+                 kgs_url='http://u-go.net/gamerecords/',
+                 index_page='train_data/kgs_index.html',
+                 data_directory='data'):
+        """Create an index of zip files containing SGF data of actual Go Games on KGS.
+
+        Parameters:
+        -----------
+        kgs_url: URL with links to zip files of games
+        index_page: Name of local html file of kgs_url
+        data_directory: name of directory relative to current path to store SGF data
+        """
+        self.kgs_url = kgs_url
+        self.index_page = index_page
+        self.data_directory = data_directory
+        self.file_info = []
+        self.urls = []
+        self.load_index()  # Load index on creation
+
+    def download_files(self):
+        """Download zip files by distributing work on all available CPUs"""
+        if not os.path.isdir(self.data_directory):
+            os.makedirs(self.data_directory)
+
+        urls_to_download = []
+        for file_info in self.file_info:
+            url = file_info['url']
+            file_name = file_info['filename']
+            if not os.path.isfile(self.data_directory + '/' + file_name):
+                urls_to_download.append((url, self.data_directory + '/' + file_name))
+        cores = multiprocessing.cpu_count()
+        pool = multiprocessing.Pool(processes=cores)
+        try:
+            it = pool.imap(worker, urls_to_download)
+            for _ in it:
+                pass
+            pool.close()
+            pool.join()
+        except KeyboardInterrupt:
+            print(">>> Caught KeyboardInterrupt, terminating workers")
+            pool.terminate()
+            pool.join()
+            sys.exit(-1)
+
+    def create_index_page(self):
+        """If there is no local html containing links to files, create one."""
+        if os.path.isfile(self.index_page):
+            print('>>> Reading cached index page')
+            index_file = open(self.index_page, 'r')
+            index_contents = index_file.read()
+            index_file.close()
+        else:
+            print('>>> Downloading index page')
+            fp = urlopen(self.kgs_url)
+            data = six.text_type(fp.read())
+            fp.close()
+            index_contents = data
+            index_file = open(self.index_page, 'w')
+            index_file.write(index_contents)
+            index_file.close()
+        return index_contents
+
+    def load_index(self):
+        """Create the actual index representation from the previously downloaded or cached html."""
+        index_contents = self.create_index_page()
+        split_page = [item for item in index_contents.split('<a href="') if item.startswith("https://")]
+        for item in split_page:
+            download_url = item.split('">Download')[0]
+            if download_url.endswith('.tar.gz'):
+                self.urls.append(download_url)
+        for url in self.urls:
+            filename = os.path.basename(url)
+            split_file_name = filename.split('-')
+            num_games = int(split_file_name[len(split_file_name) - 2])
+            print(filename + ' ' + str(num_games))
+            self.file_info.append({'url': url, 'filename': filename, 'num_games': num_games})
+
+
+if __name__ == '__main__':
+    index = KGSIndex()
+    index.download_files()

data/parallel_processor.py

@@ -0,0 +1,237 @@
+# from __future__ import print_function
+# from __future__ import absolute_import
+import os
+import glob
+import os.path
+import tarfile
+import gzip
+import shutil
+import numpy as np
+import multiprocessing
+from os import sys
+import torch
+
+from tugo.gosgf import Sgf_game
+from tugo.goboard_fast import Board, GameState, Move
+from tugo.gotypes import Player, Point
+from tugo.data.index_processor import KGSIndex
+from tugo.data.sampling import Sampler
+from tugo.data.generator import DataGenerator
+from tugo.encoders.base import get_encoder_by_name
+from torch.utils.data import TensorDataset
+
+
+def worker(jobinfo):
+    try:
+        clazz, encoder, zip_file, data_file_name, game_list = jobinfo
+        clazz(encoder=encoder).process_zip(zip_file, data_file_name, game_list)
+    except (KeyboardInterrupt, SystemExit):
+        raise Exception('>>> Exiting child process.')
+
+
+class GoDataProcessor:
+    def __init__(self, encoder='simple', data_directory='train_data'):
+        self.encoder_string = encoder
+        self.encoder = get_encoder_by_name(encoder, 19)
+        self.data_dir = data_directory
+
+    def load_go_data(self, data_type='train', num_samples=1000,
+                     use_generator=False):
+        index = KGSIndex(data_directory=self.data_dir)
+        index.download_files()
+
+        sampler = Sampler(data_dir=self.data_dir)
+        data = sampler.draw_data(data_type, num_samples)
+
+        self.map_to_workers(data_type, data)
+        if use_generator:
+            generator = DataGenerator(self.data_dir, data)
+            return generator
+        else:
+            features_and_labels = self.consolidate_games(data_type, data)
+            return features_and_labels
+
+    def unzip_data(self, zip_file_name):
+        this_gz = gzip.open(self.data_dir + '/' + zip_file_name)
+
+        tar_file = zip_file_name[0:-3]
+        this_tar = open(self.data_dir + '/' + tar_file, 'wb')
+
+        shutil.copyfileobj(this_gz, this_tar)
+        this_tar.close()
+        return tar_file
+
+    def process_zip(self, zip_file_name, data_file_name, game_list):
+        print(">>> Processing zip file:", zip_file_name, "Data file name:", data_file_name)
+        tar_file = self.unzip_data(zip_file_name)
+        zip_file = tarfile.open(self.data_dir + '/' + tar_file)
+        name_list = zip_file.getnames()
+        total_examples = self.num_total_examples(zip_file, game_list, name_list)
+        print(">>> Total examples:", total_examples)
+
+        shape = self.encoder.shape()
+        feature_shape = np.insert(shape, 0, np.asarray([total_examples]))
+        features = np.zeros(feature_shape)
+        labels = np.zeros((total_examples,))
+
+        counter = 0
+        for index in game_list:
+            name = name_list[index + 1]
+            if not name.endswith('.sgf'):
+                raise ValueError(name + ' is not a valid sgf')
+            sgf_content = zip_file.extractfile(name).read()
+            sgf = Sgf_game.from_string(sgf_content)
+            if sgf.get_handicap() is not None and sgf.get_handicap() != 0:
+                print('<func process_zip>', 'get_handicap:', sgf.get_handicap(), ' =>  skipping handicaped game ...')
+                continue
+
+            game_state, first_move_done = self.get_handicap(sgf)
+
+            for item in sgf.main_sequence_iter():
+                color, move_tuple = item.get_move()
+                point = None
+                if color is not None:
+                    if move_tuple is not None:
+                        row, col = move_tuple
+                        point = Point(row + 1, col + 1)
+                        move = Move.play(point)
+                    else:
+                        move = Move.pass_turn()
+                    if first_move_done and point is not None:
+                        features[counter] = self.encoder.encode(game_state)
+                        labels[counter] = self.encoder.encode_point(point)
+                        counter += 1
+                        game_state = game_state.apply_move(move)
+                        first_move_done = True
+
+        feature_file_base = self.data_dir + '/' + data_file_name + '_features_%d'
+        label_file_base = self.data_dir + '/' + data_file_name + '_labels_%d'
+
+        chunk = 0  # Due to files with large content, split up after chunksize
+        chunksize = 1024
+        while features.shape[0] > 0:
+            feature_file = feature_file_base % chunk
+            label_file = label_file_base % chunk
+            chunk += 1
+            # current_features, features = features[:chunksize], features[chunksize:]
+            # current_labels, labels = labels[:chunksize], labels[chunksize:]
+            current_chunksize = min(chunksize, features.shape[0])
+            current_features, features = features[:current_chunksize], features[current_chunksize:]
+            current_labels, labels = labels[:current_chunksize], labels[current_chunksize:]
+            np.save(feature_file, current_features)
+            # print(">>>>>> save feature_file:", feature_file)
+            np.save(label_file, current_labels)
+
+    def consolidate_games(self, name, samples):
+        files_needed = set(file_name for file_name, index in samples)
+        file_names = []
+        for zip_file_name in files_needed:
+            file_name = zip_file_name.replace('.tar.gz', '') + name
+            file_names.append(file_name)
+
+        feature_list = []
+        label_list = []
+        for file_name in file_names:
+            file_prefix = file_name.replace('.tar.gz', '')
+            base = self.data_dir + '/' + file_prefix + '_features_*.npy'
+            # print("base:", base)
+            # print(glob.glob(base))
+            for feature_file in glob.glob(base):
+                # print(f"Processing feature file: {feature_file}")
+                label_file = feature_file.replace('features', 'labels')
+                x = np.load(feature_file)
+                y = np.load(label_file)
+                x = x.astype('float32')
+                # y = torch.nn.functional.one_hot(torch.tensor(y.astype(int)), 19 * 19)
+                y = torch.tensor(y.astype(int))
+                # print("Feature shape:", x.shape)
+                # print("Label shape:", y.shape)
+                feature_list.append(x)
+                label_list.append(y)
+
+        assert feature_list, "No features found. Please check the data files."
+
+        features = torch.tensor(np.concatenate(feature_list, axis=0)).float()
+        labels = torch.cat(label_list, axis=0)
+
+        feature_file = self.data_dir + '/' + name + '_feature.pt'
+        label_file = self.data_dir + '/' + name + '_label.pt'
+
+        torch.save(features, feature_file)
+        torch.save(labels, label_file)
+
+        dataset = TensorDataset(torch.tensor(features), labels)
+        return dataset
+
+
+
+    @staticmethod
+    def get_handicap(sgf):  # Get handicap stones
+        go_board = Board(19, 19)
+        first_move_done = False
+        move = None
+        game_state = GameState.new_game(19)
+        if sgf.get_handicap() is not None and sgf.get_handicap() != 0:
+            # print('get_handicap2:', sgf.get_handicap())
+            for setup in sgf.get_root().get_setup_stones():
+                # print("------- setup:", setup, type(setup))
+                for move in setup:
+                    # print("------- move:", move, type(move))
+                    row, col = move
+                    point = Point(row + 1, col + 1)
+                    move = Move.play(point)
+                go_board.place_stone(Player.black, Point(row + 1, col + 1))  # black gets handicap
+            first_move_done = True
+            game_state = GameState(go_board, Player.white, None, move)
+        return game_state, first_move_done
+
+    def map_to_workers(self, data_type, samples):
+        zip_names = set()
+        indices_by_zip_name = {}
+        for filename, index in samples:
+            zip_names.add(filename)
+            if filename not in indices_by_zip_name:
+                indices_by_zip_name[filename] = []
+            indices_by_zip_name[filename].append(index)
+
+        zips_to_process = []
+        for zip_name in zip_names:
+            base_name = zip_name.replace('.tar.gz', '')
+            data_file_name = base_name + data_type
+            if not os.path.isfile(self.data_dir + '/' + data_file_name):
+                zips_to_process.append((self.__class__, self.encoder_string, zip_name,
+                data_file_name, indices_by_zip_name[zip_name]))
+                cores = multiprocessing.cpu_count()  # Determine number of CPU cores and split work load among them
+        pool = multiprocessing.Pool(processes=cores)
+        p = pool.map_async(worker, zips_to_process)
+        try:
+            _ = p.get()
+        except KeyboardInterrupt:  # Caught keyboard interrupt, terminating workers
+            pool.terminate()
+            pool.join()
+            sys.exit(-1)
+
+    def num_total_examples(self, zip_file, game_list, name_list):
+        total_examples = 0
+        for index in game_list:
+            name = name_list[index + 1]
+            if name.endswith('.sgf'):
+                sgf_content = zip_file.extractfile(name).read()
+                sgf = Sgf_game.from_string(sgf_content)
+                if sgf.get_handicap() is not None and sgf.get_handicap() != 0:
+                    print('get_handicap:', sgf.get_handicap(), ' =>  skipping handicaped game ...')
+                    continue
+                game_state, first_move_done = self.get_handicap(sgf)
+
+                num_moves = 0
+                for item in sgf.main_sequence_iter():
+                    color, move = item.get_move()
+                    if color is not None:
+                        if first_move_done:
+                            num_moves += 1
+                        first_move_done = True
+                total_examples = total_examples + num_moves
+            else:
+                raise ValueError(name + ' is not a valid sgf')
+        return total_examples
+

data/sampling.py

@@ -0,0 +1,137 @@
+# This Source Code Form is subject to the terms of the Mozilla Public License,
+# v. 2.0. If a copy of the MPL was not distributed with this file, You can
+# obtain one at http://mozilla.org/MPL/2.0/.
+from __future__ import print_function
+from __future__ import absolute_import
+import os
+import random
+from tugo.data.index_processor import KGSIndex
+from six.moves import range
+
+
+class Sampler:
+    """Sample training and test data from zipped sgf files such that test data is kept stable."""
+    def __init__(self, data_dir='data', num_test_games=100, cap_year=2015, seed=1337):
+        self.data_dir = data_dir
+        self.num_test_games = num_test_games
+        self.test_games = []
+        self.train_games = []
+        self.test_folder = 'train_data/test_samples.py'
+        self.cap_year = cap_year
+
+        random.seed(seed)
+        self.compute_test_samples()
+
+    def draw_data(self, data_type, num_samples):
+        if data_type == 'test':
+            return self.test_games
+        elif data_type == 'train' and num_samples is not None:
+            return self.draw_training_samples(num_samples)
+        elif data_type == 'train' and num_samples is None:
+            return self.draw_all_training()
+        else:
+            raise ValueError(data_type + " is not a valid data type, choose from 'train' or 'test'")
+
+    def draw_samples(self, num_sample_games):
+        """Draw num_sample_games many training games from index."""
+        available_games = []
+        index = KGSIndex(data_directory=self.data_dir)
+
+        for fileinfo in index.file_info:
+            filename = fileinfo['filename']
+            year = int(filename.split('-')[1].split('_')[0])
+            if year > self.cap_year:
+                continue
+            num_games = fileinfo['num_games']
+            for i in range(num_games):
+                available_games.append((filename, i))
+        print('>>> Total number of games used: ' + str(len(available_games)))
+
+        sample_set = set()
+        while len(sample_set) < num_sample_games:
+            sample = random.choice(available_games)
+            if sample not in sample_set:
+                sample_set.add(sample)
+        print('Drawn ' + str(num_sample_games) + ' samples:')
+        return list(sample_set)
+
+    def draw_training_games(self):
+        """Get list of all non-test games, that are no later than dec 2014
+        Ignore games after cap_year to keep training data stable
+        """
+        index = KGSIndex(data_directory=self.data_dir)
+        for file_info in index.file_info:
+            filename = file_info['filename']
+            year = int(filename.split('-')[1].split('_')[0])
+            if year > self.cap_year:
+                continue
+            num_games = file_info['num_games']
+            for i in range(num_games):
+                sample = (filename, i)
+                if sample not in self.test_games:
+                    self.train_games.append(sample)
+        print('total num training games: ' + str(len(self.train_games)))
+
+    def compute_test_samples(self):
+        """If not already existing, create local file to store fixed set of test samples"""
+        if not os.path.isfile(self.test_folder):
+            test_games = self.draw_samples(self.num_test_games)
+            test_sample_file = open(self.test_folder, 'w')
+            for sample in test_games:
+                test_sample_file.write(str(sample) + "\n")
+            test_sample_file.close()
+
+        test_sample_file = open(self.test_folder, 'r')
+        sample_contents = test_sample_file.read()
+        test_sample_file.close()
+        for line in sample_contents.split('\n'):
+            if line != "":
+                (filename, index) = eval(line)
+                self.test_games.append((filename, index))
+
+    def draw_training_samples(self, num_sample_games):
+        """Draw training games, not overlapping with any of the test games."""
+        available_games = []
+        index = KGSIndex(data_directory=self.data_dir)
+        for fileinfo in index.file_info:
+            filename = fileinfo['filename']
+            year = int(filename.split('-')[1].split('_')[0])
+            if year > self.cap_year:
+                continue
+            num_games = fileinfo['num_games']
+            for i in range(num_games):
+                available_games.append((filename, i))
+        print('total num games: ' + str(len(available_games)))
+
+        sample_set = set()
+        while len(sample_set) < num_sample_games:
+            sample = random.choice(available_games)
+            if sample not in self.test_games:
+                sample_set.add(sample)
+        print('Drawn ' + str(num_sample_games) + ' samples:')
+        return list(sample_set)
+
+    def draw_all_training(self):
+        """Draw all available training games."""
+        available_games = []
+        index = KGSIndex(data_directory=self.data_dir)
+
+        for fileinfo in index.file_info:
+            filename = fileinfo['filename']
+            year = int(filename.split('-')[1].split('_')[0])
+            if year > self.cap_year:
+                continue
+            if 'num_games' in fileinfo.keys():
+                num_games = fileinfo['num_games']
+            else:
+                continue
+            for i in range(num_games):
+                available_games.append((filename, i))
+        print('total num games: ' + str(len(available_games)))
+
+        sample_set = set()
+        for sample in available_games:
+            if sample not in self.test_games:
+                sample_set.add(sample)
+        print('Drawn all samples, ie ' + str(len(sample_set)) + ' samples:')
+        return list(sample_set)