369 lines
13 KiB
Python
369 lines
13 KiB
Python
"""
|
||
This file includes the torch models. We wrap the three
|
||
models into one class for convenience.
|
||
"""
|
||
import os
|
||
|
||
import numpy as np
|
||
|
||
import torch
|
||
import onnxruntime
|
||
from onnxruntime.datasets import get_example
|
||
from torch import nn
|
||
import torch.nn.functional as F
|
||
|
||
def to_numpy(tensor):
|
||
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
|
||
|
||
class LandlordLstmModel(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lstm = nn.LSTM(432, 128, batch_first=True)
|
||
self.dense1 = nn.Linear(887 + 128, 1024)
|
||
self.dense2 = nn.Linear(1024, 1024)
|
||
self.dense3 = nn.Linear(1024, 768)
|
||
self.dense4 = nn.Linear(768, 512)
|
||
self.dense5 = nn.Linear(512, 256)
|
||
self.dense6 = nn.Linear(256, 1)
|
||
|
||
def get_onnx_params(self, device):
|
||
return {
|
||
'args': (
|
||
torch.randn(1, 5, 432, requires_grad=True, device=device),
|
||
torch.randn(1, 887, requires_grad=True, device=device),
|
||
),
|
||
'input_names': ['z_batch','x_batch'],
|
||
'output_names': ['values'],
|
||
'dynamic_axes': {
|
||
'z_batch': {
|
||
0: "batch_size"
|
||
},
|
||
'x_batch': {
|
||
0: "batch_size"
|
||
},
|
||
'values': {
|
||
0: "batch_size"
|
||
}
|
||
}
|
||
}
|
||
|
||
def forward(self, z, x):
|
||
lstm_out, (h_n, _) = self.lstm(z)
|
||
lstm_out = lstm_out[:,-1,:]
|
||
x = torch.cat([lstm_out,x], dim=-1)
|
||
x = self.dense1(x)
|
||
x = torch.relu(x)
|
||
x = self.dense2(x)
|
||
x = torch.relu(x)
|
||
x = self.dense3(x)
|
||
x = torch.relu(x)
|
||
x = self.dense4(x)
|
||
x = torch.relu(x)
|
||
x = self.dense5(x)
|
||
x = torch.relu(x)
|
||
x = self.dense6(x)
|
||
return dict(values=x)
|
||
|
||
class FarmerLstmModel(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lstm = nn.LSTM(432, 128, batch_first=True)
|
||
self.dense1 = nn.Linear(1219 + 128, 1024)
|
||
self.dense2 = nn.Linear(1024, 1024)
|
||
self.dense3 = nn.Linear(1024, 768)
|
||
self.dense4 = nn.Linear(768, 512)
|
||
self.dense5 = nn.Linear(512, 256)
|
||
self.dense6 = nn.Linear(256, 1)
|
||
|
||
def get_onnx_params(self, device):
|
||
return {
|
||
'args': (
|
||
torch.randn(1, 5, 432, requires_grad=True, device=device),
|
||
torch.randn(1, 1219, requires_grad=True, device=device),
|
||
),
|
||
'input_names': ['z_batch','x_batch'],
|
||
'output_names': ['values'],
|
||
'dynamic_axes': {
|
||
'z_batch': {
|
||
0: "legal_actions"
|
||
},
|
||
'x_batch': {
|
||
0: "legal_actions"
|
||
},
|
||
'values': {
|
||
0: "batch_size"
|
||
}
|
||
}
|
||
}
|
||
|
||
def forward(self, z, x):
|
||
lstm_out, (h_n, _) = self.lstm(z)
|
||
lstm_out = lstm_out[:,-1,:]
|
||
x = torch.cat([lstm_out,x], dim=-1)
|
||
x = self.dense1(x)
|
||
x = torch.relu(x)
|
||
x = self.dense2(x)
|
||
x = torch.relu(x)
|
||
x = self.dense3(x)
|
||
x = torch.relu(x)
|
||
x = self.dense4(x)
|
||
x = torch.relu(x)
|
||
x = self.dense5(x)
|
||
x = torch.relu(x)
|
||
x = self.dense6(x)
|
||
return dict(values=x)
|
||
|
||
# 用于ResNet18和34的残差块，用的是2个3x3的卷积
|
||
class BasicBlock(nn.Module):
|
||
expansion = 1
|
||
|
||
def __init__(self, in_planes, planes, stride=1):
|
||
super(BasicBlock, self).__init__()
|
||
self.conv1 = nn.Conv1d(in_planes, planes, kernel_size=(3,),
|
||
stride=(stride,), padding=1, bias=False)
|
||
self.bn1 = nn.BatchNorm1d(planes)
|
||
self.conv2 = nn.Conv1d(planes, planes, kernel_size=(3,),
|
||
stride=(1,), padding=1, bias=False)
|
||
self.bn2 = nn.BatchNorm1d(planes)
|
||
self.shortcut = nn.Sequential()
|
||
# 经过处理后的x要与x的维度相同(尺寸和深度)
|
||
# 如果不相同，需要添加卷积+BN来变换为同一维度
|
||
if stride != 1 or in_planes != self.expansion * planes:
|
||
self.shortcut = nn.Sequential(
|
||
nn.Conv1d(in_planes, self.expansion * planes,
|
||
kernel_size=(1,), stride=(stride,), bias=False),
|
||
nn.BatchNorm1d(self.expansion * planes)
|
||
)
|
||
|
||
def forward(self, x):
|
||
out = F.relu(self.bn1(self.conv1(x)))
|
||
out = self.bn2(self.conv2(out))
|
||
out += self.shortcut(x)
|
||
out = F.relu(out)
|
||
return out
|
||
|
||
class GeneralModel(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.in_planes = 80
|
||
#input 1*108*41
|
||
self.conv1 = nn.Conv1d(40, 80, kernel_size=(3,),
|
||
stride=(2,), padding=1, bias=False) #1*108*80
|
||
|
||
self.bn1 = nn.BatchNorm1d(80)
|
||
|
||
self.layer1 = self._make_layer(BasicBlock, 80, 2, stride=2)#1*27*80
|
||
self.layer2 = self._make_layer(BasicBlock, 160, 2, stride=2)#1*14*160
|
||
self.layer3 = self._make_layer(BasicBlock, 320, 2, stride=2)#1*7*320
|
||
self.layer4 = self._make_layer(BasicBlock, 640, 2, stride=2)#1*4*640
|
||
# self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
|
||
self.linear1 = nn.Linear(640 * BasicBlock.expansion * 4 + 56, 2048)
|
||
self.linear2 = nn.Linear(2048, 1024)
|
||
self.linear3 = nn.Linear(1024, 512)
|
||
self.linear4 = nn.Linear(512, 256)
|
||
self.linear5 = nn.Linear(256, 1)
|
||
|
||
def _make_layer(self, block, planes, num_blocks, stride):
|
||
strides = [stride] + [1] * (num_blocks - 1)
|
||
layers = []
|
||
for stride in strides:
|
||
layers.append(block(self.in_planes, planes, stride))
|
||
self.in_planes = planes * block.expansion
|
||
return nn.Sequential(*layers)
|
||
|
||
def get_onnx_params(self, device=None):
|
||
return {
|
||
'args': (
|
||
torch.randn(1, 40, 108, requires_grad=True, device=device),
|
||
torch.randn(1, 56, requires_grad=True, device=device),
|
||
),
|
||
'input_names': ['z_batch','x_batch'],
|
||
'output_names': ['values'],
|
||
'dynamic_axes': {
|
||
'z_batch': {
|
||
0: "batch_size"
|
||
},
|
||
'x_batch': {
|
||
0: "batch_size"
|
||
},
|
||
'values': {
|
||
0: "batch_size"
|
||
}
|
||
}
|
||
}
|
||
|
||
def forward(self, z, x):
|
||
out = F.relu(self.bn1(self.conv1(z)))
|
||
out = self.layer1(out)
|
||
out = self.layer2(out)
|
||
out = self.layer3(out)
|
||
out = self.layer4(out)
|
||
out = out.flatten(1,2)
|
||
out = torch.cat([x,out], dim=-1)
|
||
out = F.leaky_relu_(self.linear1(out))
|
||
out = F.leaky_relu_(self.linear2(out))
|
||
out = F.leaky_relu_(self.linear3(out))
|
||
out = F.leaky_relu_(self.linear4(out))
|
||
out = F.leaky_relu_(self.linear5(out))
|
||
return dict(values=out)
|
||
|
||
|
||
# Model dict is only used in evaluation but not training
|
||
model_dict = {}
|
||
model_dict['landlord'] = LandlordLstmModel
|
||
model_dict['landlord_up'] = FarmerLstmModel
|
||
model_dict['landlord_front'] = FarmerLstmModel
|
||
model_dict['landlord_down'] = FarmerLstmModel
|
||
model_dict_new = {}
|
||
model_dict_new['landlord'] = GeneralModel
|
||
model_dict_new['landlord_up'] = GeneralModel
|
||
model_dict_new['landlord_front'] = GeneralModel
|
||
model_dict_new['landlord_down'] = GeneralModel
|
||
|
||
class OldModel:
|
||
"""
|
||
The wrapper for the three models. We also wrap several
|
||
interfaces such as share_memory, eval, etc.
|
||
"""
|
||
def __init__(self, device=0, flags=None):
|
||
self.models = {}
|
||
if not device == "cpu":
|
||
device = 'cuda:' + str(device)
|
||
self.device = torch.device(device)
|
||
self.models['landlord'] = LandlordLstmModel().to(self.device)
|
||
self.models['landlord_up'] = FarmerLstmModel().to(self.device)
|
||
self.models['landlord_front'] = FarmerLstmModel().to(self.device)
|
||
self.models['landlord_down'] = FarmerLstmModel().to(self.device)
|
||
self.onnx_models = {
|
||
'landlord': None,
|
||
'landlord_up': None,
|
||
'landlord_front': None,
|
||
'landlord_down': None
|
||
}
|
||
|
||
def set_onnx_model(self, position, model_path):
|
||
self.onnx_models[position] = onnxruntime.InferenceSession(get_example(model_path))
|
||
|
||
def get_onnx_params(self, position):
|
||
self.models[position].get_onnx_params(self.device)
|
||
|
||
def forward(self, position, z, x, return_value=False, flags=None):
|
||
model = self.onnx_models[position]
|
||
if model is None:
|
||
model = self.models[position]
|
||
values = model.forward(z, x)['values']
|
||
else:
|
||
onnx_out = model.run(None, {'z_batch': to_numpy(z), 'x_batch': to_numpy(x)})
|
||
values = torch.tensor(onnx_out[0])
|
||
if return_value:
|
||
return dict(values=values)
|
||
else:
|
||
if flags is not None and flags.exp_epsilon > 0 and np.random.rand() < flags.exp_epsilon:
|
||
action = torch.randint(values.shape[0], (1,))[0]
|
||
else:
|
||
action = torch.argmax(values,dim=0)[0]
|
||
return dict(action=action)
|
||
|
||
def share_memory(self):
|
||
if self.models['landlord'] is not None:
|
||
self.models['landlord'].share_memory()
|
||
self.models['landlord_up'].share_memory()
|
||
self.models['landlord_front'].share_memory()
|
||
self.models['landlord_down'].share_memory()
|
||
|
||
def eval(self):
|
||
if self.models['landlord'] is not None:
|
||
self.models['landlord'].eval()
|
||
self.models['landlord_up'].eval()
|
||
self.models['landlord_front'].eval()
|
||
self.models['landlord_down'].eval()
|
||
|
||
def parameters(self, position):
|
||
return self.models[position].parameters()
|
||
|
||
def get_model(self, position):
|
||
return self.models[position]
|
||
|
||
def get_models(self):
|
||
return self.models
|
||
|
||
|
||
class Model:
|
||
"""
|
||
The wrapper for the three models. We also wrap several
|
||
interfaces such as share_memory, eval, etc.
|
||
"""
|
||
def __init__(self, device=0, flags=None):
|
||
self.models = {}
|
||
self.onnx_models = {}
|
||
self.flags = flags
|
||
if not device == "cpu":
|
||
device = 'cuda:' + str(device)
|
||
self.device = torch.device(device)
|
||
# model = GeneralModel().to(torch.device(device))
|
||
positions = ['landlord', 'landlord_up', 'landlord_front', 'landlord_down']
|
||
if flags is not None and flags.enable_onnx:
|
||
for position in positions:
|
||
self.models[position] = None
|
||
else:
|
||
for position in positions:
|
||
self.models[position] = GeneralModel().to(self.device)
|
||
self.onnx_models = {
|
||
'landlord': None,
|
||
'landlord_up': None,
|
||
'landlord_front': None,
|
||
'landlord_down': None
|
||
}
|
||
|
||
def set_onnx_model(self):
|
||
positions = ['landlord', 'landlord_up', 'landlord_front', 'landlord_down']
|
||
for position in positions:
|
||
model_path = os.path.abspath('%s/%s/model_%s.onnx' % (self.flags.savedir, self.flags.xpid, position))
|
||
self.onnx_models[position] = onnxruntime.InferenceSession(get_example(model_path), providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
||
|
||
def get_onnx_params(self, position):
|
||
self.models[position].get_onnx_params(self.device)
|
||
|
||
def forward(self, position, z, x, return_value=False, flags=None, debug=False):
|
||
if self.flags.enable_onnx and len(self.onnx_models) == 0:
|
||
self.set_onnx_model()
|
||
model = self.onnx_models[position]
|
||
if model is None:
|
||
model = self.models[position]
|
||
values = model.forward(z, x)['values']
|
||
else:
|
||
onnx_out = model.run(None, {'z_batch': to_numpy(z), 'x_batch': to_numpy(x)})
|
||
values = torch.tensor(onnx_out[0])
|
||
if return_value:
|
||
return dict(values=values)
|
||
else:
|
||
if flags is not None and flags.exp_epsilon > 0 and np.random.rand() < flags.exp_epsilon:
|
||
action = torch.randint(values.shape[0], (1,))[0]
|
||
else:
|
||
action = torch.argmax(values,dim=0)[0]
|
||
return dict(action=action)
|
||
|
||
def share_memory(self):
|
||
if self.models['landlord'] is not None:
|
||
self.models['landlord'].share_memory()
|
||
self.models['landlord_up'].share_memory()
|
||
self.models['landlord_front'].share_memory()
|
||
self.models['landlord_down'].share_memory()
|
||
|
||
def eval(self):
|
||
if self.models['landlord'] is not None:
|
||
self.models['landlord'].eval()
|
||
self.models['landlord_up'].eval()
|
||
self.models['landlord_front'].eval()
|
||
self.models['landlord_down'].eval()
|
||
|
||
def parameters(self, position):
|
||
return self.models[position].parameters()
|
||
|
||
def get_model(self, position):
|
||
return self.models[position]
|
||
|
||
def get_models(self):
|
||
return self.models
|
||
|