# 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('', '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