Spaces:
Sleeping
Sleeping
File size: 5,068 Bytes
0154f58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
import pickle
from pathlib import Path
from typing import Dict, Type, Union, Literal, List
import torch
from my_chess.learner.models import (
Model,
ModelConfig,
DeepChessAlphaBeta,
DeepChessAlphaBetaConfig,
DeepChessEvaluator,
DeepChessEvaluatorConfig,
DeepChessFE,
DeepChessFEConfig
)
from huggingface_hub import PyTorchModelHubMixin
class DCFE(DeepChessFE, PyTorchModelHubMixin):
def __init__(
self,
hidden_dims:Union[int, List[int]]=[4096, 1024, 256, 128],
activations:Union[str, List[str]]='relu',
batch_norm:bool = True) -> None:
input_sample = torch.rand(1,8,8,111)
super().__init__(input_sample, config=DeepChessFEConfig(**dict(
hidden_dims=hidden_dims,
activations=activations,
batch_norm=batch_norm,
)))
class DCEvaluator(DeepChessEvaluator, PyTorchModelHubMixin):
def __init__(
self,
feature_extractor:Type[Model]=None,
feature_extractor_config:ModelConfig=None,
feature_extractor_param_dir:Union[str, Path]=None,
hidden_dims:Union[int, List[int]]=[512, 252, 128],
activations:Union[str, List[str]]='relu',
batch_norm:bool=True,) -> None:
input_sample = torch.rand(1,8,8,111)
if feature_extractor is None:
feature_extractor = DCFE.from_pretrained("mzimm003/DeepChessReplicationFeatureExtractor")
super().__init__(input_sample, config=DeepChessEvaluatorConfig(**dict(
feature_extractor=feature_extractor,
feature_extractor_config=feature_extractor_config,
feature_extractor_param_dir=feature_extractor_param_dir,
hidden_dims=hidden_dims,
activations=activations,
batch_norm=batch_norm,
)))
class DCMinMax(DeepChessAlphaBeta, PyTorchModelHubMixin):
def __init__(
self,
board_evaluator:Type[Model]=None,
board_evaluator_config:ModelConfig=None,
board_evaluator_param_dir:Union[str, Path]=None,
max_depth:int = 8,
iterate_depths:bool = True,
move_sort:Literal['none', 'random', 'evaluation'] = 'evaluation') -> None:
input_sample = torch.rand(1,8,8,111)
if board_evaluator is None:
board_evaluator = DCFE.from_pretrained("mzimm003/DeepChessReplicationBoardEvaluator")
super().__init__(input_sample, config=DeepChessAlphaBetaConfig(**dict(
board_evaluator=board_evaluator,
board_evaluator_config=board_evaluator_config,
board_evaluator_param_dir=board_evaluator_param_dir,
max_depth=max_depth,
iterate_depths=iterate_depths,
move_sort=move_sort,
)))
def get_model_attrs(dir:Union[str, Path]):
best_model_dir = Path(dir).resolve()
best_model_class = None
best_model_config = None
with open(best_model_dir/"params.pkl",'rb') as f:
x = pickle.load(f)
best_model_class = x['model']
best_model_config = x['model_config']
latest_checkpoint = sorted(best_model_dir.glob('checkpoint*'), reverse=True)[0]/'model.pt'
return best_model_class, best_model_config, latest_checkpoint
def main1(kwargs=None):
mod_cls, mod_config, mod_chkpt = get_model_attrs("/opt/ray/results/ChessFeatureExtractor/AutoEncoder_1557d_00000_0_batch_size=256,model_config=ref_ph_a52f5213,lr=0.0001_2024-07-12_22-30-58")
model = DCFE(**mod_config.asDict())
model.load_state_dict(torch.load(mod_chkpt,map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")))
model.save_pretrained('/opt/pretrained_models/DCFE')
model.push_to_hub(
"mzimm003/DeepChessReplicationFeatureExtractor")
mod_cls, mod_config, mod_chkpt = get_model_attrs("/opt/ray/results/DeepChessEvaluator/ChessEvaluation_9866d_00000_0_learning_rate_scheduler_config=step_size_200_gamma_0_9,model_config=ref_ph_a52f5213,lr=0.1000_2024-07-13_09-18-52")
model = DCEvaluator(**mod_config.asDict())
model.load_state_dict(torch.load(mod_chkpt,map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")))
model.save_pretrained('/opt/pretrained_models/DCEvaluator')
model.push_to_hub(
"mzimm003/DeepChessReplicationBoardEvaluator")
model = DCMinMax(
**dict(
board_evaluator=mod_cls,
board_evaluator_config=mod_config,
board_evaluator_param_dir=mod_chkpt,
max_depth=3,
move_sort='evaluation'))
model.save_pretrained('/opt/pretrained_models/DCMinMax')
model.push_to_hub(
"mzimm003/DeepChessReplicationMinMax")
def main(kwargs=None):
modelfe = DCFE.from_pretrained("mzimm003/DeepChessReplicationFeatureExtractor")
modeleval = DCEvaluator.from_pretrained("mzimm003/DeepChessReplicationBoardEvaluator")
modelminmax = DCMinMax.from_pretrained("mzimm003/DeepChessReplicationMinMax")
pass
if __name__ == "__main__":
main() |