liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-16 09:11:00 +08:00
Code Activity

update format

Browse Source
This commit is contained in:
Alex Wang
2021-01-21 22:18:40 +09:00
committed by you-n-g
parent 5da5ad4b9f
commit 0a86a6f392
1 changed files with 118 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
qlib/contrib/model/pytorch_tabnet.py
View File

@@ -28,12 +28,34 @@ from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP from ...data.dataset.handler import DataHandlerLP
class TabNet_Model(Model): class TabNet_Model(Model):
def __init__(self, d_feat=158, out_dim = 64, final_out_dim = 1, batch_size = 4096, n_d=64, n_a=64, n_shared=2, n_ind=2, def __init__(
n_steps=5, n_epochs=100, pretrain_n_epochs=50, relax=1.3, vbs=2048, seed = 993, optimizer='adam', loss = 'mse', self,
metric = '', early_stop = 20, GPU='1', pretrain_loss = 'custom', ps = 0.3, lr = 0.01, pretrain = True, pretrain_file = './pretrain/best.model'): d_feat=158,
out_dim=64,
final_out_dim=1,
batch_size=4096,
n_d=64,
n_a=64,
n_shared=2,
n_ind=2,
n_steps=5,
n_epochs=100,
pretrain_n_epochs=50,
relax=1.3,
vbs=2048,
seed=993,
optimizer="adam",
loss="mse",
metric="",
early_stop=20,
GPU="1",
pretrain_loss="custom",
ps=0.3,
lr=0.01,
pretrain=True,
pretrain_file="./pretrain/best.model",
):
""" """
TabNet model for Qlib TabNet model for Qlib
@@ -64,35 +86,37 @@ class TabNet_Model(Model):
"\nbatch_size : {}" "\nbatch_size : {}"
"\nvirtual bs : {}" "\nvirtual bs : {}"
"\nGPU : {}" "\nGPU : {}"
"\npretrain: {}".format( "\npretrain: {}".format(self.batch_size, vbs, GPU, pretrain)
self.batch_size,
vbs,
GPU,
pretrain
)
) )
np.random.seed(self.seed) np.random.seed(self.seed)
torch.manual_seed(self.seed) torch.manual_seed(self.seed)
self.tabnet_model = TabNet(inp_dim = self.d_feat, out_dim = self.out_dim, vbs = vbs, relax = relax, device = self.device).to(self.device) self.tabnet_model = TabNet(
self.tabnet_decoder = TabNet_Decoder(self.out_dim, self.d_feat, n_shared, n_ind, vbs, n_steps, self.device).to(self.device) inp_dim=self.d_feat, out_dim=self.out_dim, vbs=vbs, relax=relax, device=self.device
).to(self.device)
self.tabnet_decoder = TabNet_Decoder(self.out_dim, self.d_feat, n_shared, n_ind, vbs, n_steps, self.device).to(
self.device
)
if optimizer.lower() == "adam": if optimizer.lower() == "adam":
self.pretrain_optimizer = optim.Adam(list(self.tabnet_model.parameters())+list(self.tabnet_decoder.parameters()), lr=self.lr) self.pretrain_optimizer = optim.Adam(
list(self.tabnet_model.parameters()) + list(self.tabnet_decoder.parameters()), lr=self.lr
)
self.train_optimizer = optim.Adam(self.tabnet_model.parameters(), lr=self.lr) self.train_optimizer = optim.Adam(self.tabnet_model.parameters(), lr=self.lr)
elif optimizer.lower() == "gd": elif optimizer.lower() == "gd":
self.pretrain_optimizer = optim.SGD(list(self.tabnet_model.parameters())+list(self.tabnet_decoder.parameters()), lr=self.lr) self.pretrain_optimizer = optim.SGD(
list(self.tabnet_model.parameters()) + list(self.tabnet_decoder.parameters()), lr=self.lr
)
self.train_optimizer = optim.SGD(self.tabnet_model.parameters(), lr=self.lr) self.train_optimizer = optim.SGD(self.tabnet_model.parameters(), lr=self.lr)
else: else:
raise NotImplementedError("optimizer {} is not supported!".format(optimizer)) raise NotImplementedError("optimizer {} is not supported!".format(optimizer))
def pretrain_fn(self, dataset=DatasetH, pretrain_file="./pretrain/best.model"):
def pretrain_fn(self, dataset = DatasetH, pretrain_file = './pretrain/best.model'):
# make a directory if pretrian director does not exist # make a directory if pretrian director does not exist
if pretrain_file.startswith('./pretrain') and not os.path.exists('pretrain'): if pretrain_file.startswith("./pretrain") and not os.path.exists("pretrain"):
self.logger.info("make folder to store model...") self.logger.info("make folder to store model...")
os.makedirs('pretrain') os.makedirs("pretrain")
[df_train, df_valid] = dataset.prepare( [df_train, df_valid] = dataset.prepare(
["pretrain", "pretrain_validation"], ["pretrain", "pretrain_validation"],
@@ -112,7 +136,7 @@ class TabNet_Model(Model):
best_loss = np.inf best_loss = np.inf
for epoch_idx in range(self.pretrain_n_epochs): for epoch_idx in range(self.pretrain_n_epochs):
self.logger.info('epoch: %s' % (epoch_idx)) self.logger.info("epoch: %s" % (epoch_idx))
self.logger.info("pre-training...") self.logger.info("pre-training...")
self.pretrain_epoch(x_train) self.pretrain_epoch(x_train)
self.logger.info("evaluating...") self.logger.info("evaluating...")
@@ -130,7 +154,6 @@ class TabNet_Model(Model):
self.logger.info("early stop") self.logger.info("early stop")
break break
def fit( def fit(
self, self,
dataset: DatasetH, dataset: DatasetH,
@@ -138,7 +161,7 @@ class TabNet_Model(Model):
verbose=True, verbose=True,
save_path=None, save_path=None,
): ):
if(self.pretrain): if self.pretrain:
# there is a pretrained model, load the model # there is a pretrained model, load the model
self.logger.info("Pretrain...") self.logger.info("Pretrain...")
self.pretrain_fn(dataset, self.pretrain_file) self.pretrain_fn(dataset, self.pretrain_file)
@@ -167,7 +190,7 @@ class TabNet_Model(Model):
self._fitted = True self._fitted = True
for epoch_idx in range(self.n_epochs): for epoch_idx in range(self.n_epochs):
self.logger.info('epoch: %s' % (epoch_idx)) self.logger.info("epoch: %s" % (epoch_idx))
self.logger.info("training...") self.logger.info("training...")
self.train_epoch(x_train, y_train) self.train_epoch(x_train, y_train)
self.logger.info("evaluating...") self.logger.info("evaluating...")
@@ -188,7 +211,6 @@ class TabNet_Model(Model):
break break
self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch)) self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
def predict(self, dataset): def predict(self, dataset):
if not self._fitted: if not self._fitted:
raise ValueError("model is not fitted yet!") raise ValueError("model is not fitted yet!")
@@ -217,7 +239,6 @@ class TabNet_Model(Model):
return pd.Series(np.concatenate(preds), index=index) return pd.Series(np.concatenate(preds), index=index)
def test_epoch(self, data_x, data_y): def test_epoch(self, data_x, data_y):
# prepare training data # prepare training data
x_values = torch.from_numpy(data_x.values) x_values = torch.from_numpy(data_x.values)
@@ -364,10 +385,12 @@ class FinetuneModel(nn.Module):
""" """
FinuetuneModel for adding a layer by the end FinuetuneModel for adding a layer by the end
""" """
def __init__(self, input_dim, output_dim, trained_model): def __init__(self, input_dim, output_dim, trained_model):
super().__init__() super().__init__()
self.model = trained_model self.model = trained_model
self.fc = nn.Linear(input_dim, output_dim) self.fc = nn.Linear(input_dim, output_dim)
def forward(self, x, priors): def forward(self, x, priors):
return self.fc(self.model(x, priors)[0]).squeeze() # take the vec out return self.fc(self.model(x, priors)[0]).squeeze() # take the vec out
@@ -378,13 +401,11 @@ class DecoderStep(nn.Module):
self.fea_tran = FeatureTransformer(inp_dim, out_dim, shared, n_ind, vbs, device) self.fea_tran = FeatureTransformer(inp_dim, out_dim, shared, n_ind, vbs, device)
self.fc = nn.Linear(out_dim, out_dim) self.fc = nn.Linear(out_dim, out_dim)
def forward(self, x): def forward(self, x):
x = self.fea_tran(x) x = self.fea_tran(x)
return self.fc(x) return self.fc(x)
class TabNet_Decoder(nn.Module): class TabNet_Decoder(nn.Module):
def __init__(self, inp_dim, out_dim, n_shared, n_ind, vbs, n_steps, device): def __init__(self, inp_dim, out_dim, n_shared, n_ind, vbs, n_steps, device):
""" """
@@ -412,9 +433,10 @@ class TabNet_Decoder(nn.Module):
return out return out
class TabNet(nn.Module): class TabNet(nn.Module):
def __init__(self, inp_dim=6, out_dim=6, n_d=64, n_a=64, n_shared=2, n_ind=2, n_steps=5, relax=1.2, vbs=1024, device = 'cpu'): def __init__(
self, inp_dim=6, out_dim=6, n_d=64, n_a=64, n_shared=2, n_ind=2, n_steps=5, relax=1.2, vbs=1024, device="cpu"
):
""" """
TabNet AKA the original encoder TabNet AKA the original encoder
@@ -468,6 +490,7 @@ class GBN(nn.Module):
Args: Args:
vbs: virtual batch size vbs: virtual batch size
""" """
def __init__(self, inp, vbs=1024, momentum=0.01): def __init__(self, inp, vbs=1024, momentum=0.01):
super().__init__() super().__init__()
self.bn = nn.BatchNorm1d(inp, momentum=momentum) self.bn = nn.BatchNorm1d(inp, momentum=momentum)
@@ -486,6 +509,7 @@ class GLU(nn.Module):
Args: Args:
vbs: virtual batch size vbs: virtual batch size
""" """
def __init__(self, inp_dim, out_dim, fc=None, vbs=1024): def __init__(self, inp_dim, out_dim, fc=None, vbs=1024):
super().__init__() super().__init__()
if fc: if fc:
@@ -507,6 +531,7 @@ class AttentionTransformer(nn.Module):
use the same features more. When it is set to 1 use the same features more. When it is set to 1
we can use every feature only once we can use every feature only once
""" """
def __init__(self, d_a, inp_dim, relax, vbs=1024): def __init__(self, d_a, inp_dim, relax, vbs=1024):
super().__init__() super().__init__()
self.fc = nn.Linear(d_a, inp_dim) self.fc = nn.Linear(d_a, inp_dim)
@@ -538,7 +563,7 @@ class FeatureTransformer(nn.Module):
self.independ.append(GLU(inp, out_dim, vbs=vbs)) self.independ.append(GLU(inp, out_dim, vbs=vbs))
for x in range(first, n_ind): for x in range(first, n_ind):
self.independ.append(GLU(out_dim, out_dim, vbs=vbs)) self.independ.append(GLU(out_dim, out_dim, vbs=vbs))
self.scale = torch.sqrt(torch.tensor([.5], device=device)) self.scale = torch.sqrt(torch.tensor([0.5], device=device))
def forward(self, x): def forward(self, x):
if self.shared: if self.shared:
@@ -552,11 +577,11 @@ class FeatureTransformer(nn.Module):
return x return x
class DecisionStep(nn.Module): class DecisionStep(nn.Module):
""" """
One step for the TabNet One step for the TabNet
""" """
def __init__(self, inp_dim, n_d, n_a, shared, n_ind, relax, vbs, device): def __init__(self, inp_dim, n_d, n_a, shared, n_ind, relax, vbs, device):
super().__init__() super().__init__()
self.atten_tran = AttentionTransformer(n_a, inp_dim, relax, vbs) self.atten_tran = AttentionTransformer(n_a, inp_dim, relax, vbs)
@@ -569,7 +594,6 @@ class DecisionStep(nn.Module):
return x, sparse_loss return x, sparse_loss
def make_ix_like(input, dim=0): def make_ix_like(input, dim=0):
d = input.size(dim) d = input.size(dim)
rho = torch.arange(1, d + 1, device=input.device, dtype=input.dtype) rho = torch.arange(1, d + 1, device=input.device, dtype=input.dtype)
@@ -577,10 +601,12 @@ def make_ix_like(input, dim=0):
view[0] = -1 view[0] = -1
return rho.view(view).transpose(0, dim) return rho.view(view).transpose(0, dim)
class SparsemaxFunction(Function): class SparsemaxFunction(Function):
""" """
SparseMax function for replacing reLU SparseMax function for replacing reLU
""" """
@staticmethod @staticmethod
def forward(ctx, input, dim=-1): def forward(ctx, input, dim=-1):
ctx.dim = dim ctx.dim = dim