bbig/tugo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update directory structure

Browse Source
This commit is contained in:
bbig
2023-05-30 17:16:48 +08:00
parent 4968e6400e
commit 50111160e3
37 changed files with 150 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data_processing/__init__.py
View File
data_processing/generator.py
+40
View File
@@ -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_processing/index_processor.py
+107
View File
@@ -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_processing/parallel_processor.py
+237
View File
@@ -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 game_logic.goboard_fast import Board, GameState, Move
from game_logic.gotypes import Player, Point
from tugo.data_processing.index_processor import KGSIndex
from tugo.data_processing.sampling import Sampler
from tugo.data_processing.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_processing/sampling.py
+137
View File
@@ -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_processing.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)