87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
import torch
|
|
import numpy as np
|
|
import os
|
|
import onnxruntime
|
|
from onnxruntime.datasets import get_example
|
|
|
|
from douzero.env.env import get_obs
|
|
|
|
def _load_model(position, model_path, model_type, use_legacy):
|
|
from douzero.dmc.models import model_dict_new, model_dict, model_dict_legacy
|
|
model = None
|
|
if model_type == "general":
|
|
model = model_dict_new[position]()
|
|
else:
|
|
if use_legacy:
|
|
model = model_dict_legacy[position]()
|
|
else:
|
|
model = model_dict[position]()
|
|
model_state_dict = model.state_dict()
|
|
if torch.cuda.is_available():
|
|
pretrained = torch.load(model_path, map_location='cuda:0')
|
|
else:
|
|
pretrained = torch.load(model_path, map_location='cpu')
|
|
pretrained = {k: v for k, v in pretrained.items() if k in model_state_dict}
|
|
model_state_dict.update(pretrained)
|
|
model.load_state_dict(model_state_dict)
|
|
# torch.save(model.state_dict(), model_path.replace(".ckpt", "_nobn.ckpt"))
|
|
if torch.cuda.is_available():
|
|
model.cuda()
|
|
model.eval()
|
|
onnx_params = model.get_onnx_params(torch.device('cpu'))
|
|
model_path = model_path + '.onnx'
|
|
if not os.path.exists(model_path):
|
|
torch.onnx.export(
|
|
model,
|
|
onnx_params['args'],
|
|
model_path,
|
|
export_params=True,
|
|
opset_version=10,
|
|
do_constant_folding=True,
|
|
input_names=onnx_params['input_names'],
|
|
output_names=onnx_params['output_names'],
|
|
dynamic_axes=onnx_params['dynamic_axes']
|
|
)
|
|
|
|
return model
|
|
|
|
class DeepAgent:
|
|
|
|
def __init__(self, position, model_path):
|
|
self.use_legacy = True if "legacy" in model_path else False
|
|
self.model_type = "general" if "resnet" in model_path else "old"
|
|
self.model = _load_model(position, model_path, self.model_type, self.use_legacy)
|
|
self.onnx_model = onnxruntime.InferenceSession(get_example(os.path.abspath(model_path + '.onnx')), providers=['CPUExecutionProvider'])
|
|
self.EnvCard2RealCard = {3: '3', 4: '4', 5: '5', 6: '6', 7: '7',
|
|
8: '8', 9: '9', 10: 'T', 11: 'J', 12: 'Q',
|
|
13: 'K', 14: 'A', 17: '2', 20: 'X', 30: 'D'}
|
|
def act(self, infoset):
|
|
if len(infoset.legal_actions) == 1:
|
|
return infoset.legal_actions[0]
|
|
|
|
obs = get_obs(infoset, self.model_type == "general", self.use_legacy)
|
|
|
|
z_batch = torch.from_numpy(obs['z_batch']).float()
|
|
x_batch = torch.from_numpy(obs['x_batch']).float()
|
|
if torch.cuda.is_available():
|
|
z_batch, x_batch = z_batch.cuda(), x_batch.cuda()
|
|
# y_pred = self.model.forward(z_batch, x_batch)['values']
|
|
# y_pred = y_pred.detach().cpu().numpy()
|
|
y_pred = self.onnx_model.run(None, {'z_batch': obs['z_batch'], 'x_batch': obs['x_batch']})[0]
|
|
|
|
best_action_index = np.argmax(y_pred, axis=0)[0]
|
|
best_action = infoset.legal_actions[best_action_index]
|
|
# action_list = []
|
|
# output = ""
|
|
# for i, action in enumerate(y_pred):
|
|
# action_list.append((y_pred[i].item(), "".join([self.EnvCard2RealCard[ii] for ii in infoset.legal_actions[i]]) if len(infoset.legal_actions[i]) != 0 else "Pass"))
|
|
# action_list.sort(key=lambda x: x[0], reverse=True)
|
|
# value_list = []
|
|
# for action in action_list:
|
|
# output += str(round(action[0],3)) + " " + action[1] + "\n"
|
|
# value_list.append(action[0])
|
|
# # print(value_list)
|
|
# print(output)
|
|
# print("--------------------\n")
|
|
return best_action
|