liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-14 16:26:55 +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

168
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
@@ -60,39 +82,41 @@ class TabNet_Model(Model):
self.pretrain = pretrain self.pretrain = pretrain
self.pretrain_file = pretrain_file self.pretrain_file = pretrain_file
self.logger.info( self.logger.info(
"TabNet:" "TabNet:"
"\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"],
@@ -100,8 +124,8 @@ class TabNet_Model(Model):
data_key=DataHandlerLP.DK_L, data_key=DataHandlerLP.DK_L,
) )
df_train.fillna(df_train.mean(), inplace = True) df_train.fillna(df_train.mean(), inplace=True)
df_valid.fillna(df_valid.mean(), inplace = True) df_valid.fillna(df_valid.mean(), inplace=True)
x_train = df_train["feature"] x_train = df_train["feature"]
x_valid = df_valid["feature"] x_valid = df_valid["feature"]
@@ -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...")
@@ -125,12 +149,11 @@ class TabNet_Model(Model):
torch.save(self.tabnet_model.state_dict(), pretrain_file) torch.save(self.tabnet_model.state_dict(), pretrain_file)
best_loss = valid_loss best_loss = valid_loss
else: else:
stop_steps+=1 stop_steps += 1
if stop_steps >= self.early_stop: if stop_steps >= self.early_stop:
self.logger.info("early stop") self.logger.info("early stop")
break break
def fit( def fit(
self, self,
dataset: DatasetH, dataset: DatasetH,
@@ -138,21 +161,21 @@ 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)
self.logger.info("Load Pretrain model") self.logger.info("Load Pretrain model")
self.tabnet_model.load_state_dict(torch.load(self.pretrain_file)) self.tabnet_model.load_state_dict(torch.load(self.pretrain_file))
#adding one more linear layer to fit the final output dimension # adding one more linear layer to fit the final output dimension
self.tabnet_model = FinetuneModel(self.out_dim, self.final_out_dim, self.tabnet_model).to(self.device) self.tabnet_model = FinetuneModel(self.out_dim, self.final_out_dim, self.tabnet_model).to(self.device)
df_train, df_valid = dataset.prepare( df_train, df_valid = dataset.prepare(
["train", "valid"], ["train", "valid"],
col_set=["feature", "label"], col_set=["feature", "label"],
data_key=DataHandlerLP.DK_L, data_key=DataHandlerLP.DK_L,
) )
df_train.fillna(df_train.mean(), inplace = True) df_train.fillna(df_train.mean(), inplace=True)
x_train, y_train = df_train["feature"], df_train["label"] x_train, y_train = df_train["feature"], df_train["label"]
x_valid, y_valid = df_valid["feature"], df_valid["label"] x_valid, y_valid = df_valid["feature"], df_valid["label"]
@@ -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)
@@ -287,8 +308,8 @@ class TabNet_Model(Model):
if len(indices) - i < self.batch_size: if len(indices) - i < self.batch_size:
break break
S_mask = torch.bernoulli(torch.empty(self.batch_size, self.d_feat).fill_(self.ps)) S_mask = torch.bernoulli(torch.empty(self.batch_size, self.d_feat).fill_(self.ps))
x_train_values = train_set[indices[i : i + self.batch_size]] * (1-S_mask) x_train_values = train_set[indices[i : i + self.batch_size]] * (1 - S_mask)
y_train_values = train_set[indices[i : i + self.batch_size]] * (S_mask) y_train_values = train_set[indices[i : i + self.batch_size]] * (S_mask)
S_mask = S_mask.to(self.device) S_mask = S_mask.to(self.device)
@@ -318,14 +339,14 @@ class TabNet_Model(Model):
if len(indices) - i < self.batch_size: if len(indices) - i < self.batch_size:
break break
S_mask = torch.bernoulli(torch.empty(self.batch_size, self.d_feat).fill_(self.ps)) S_mask = torch.bernoulli(torch.empty(self.batch_size, self.d_feat).fill_(self.ps))
x_train_values = train_set[indices[i : i + self.batch_size]] * (1-S_mask) x_train_values = train_set[indices[i : i + self.batch_size]] * (1 - S_mask)
y_train_values = train_set[indices[i : i + self.batch_size]] * (S_mask) y_train_values = train_set[indices[i : i + self.batch_size]] * (S_mask)
feature = x_train_values.float().to(self.device) feature = x_train_values.float().to(self.device)
label = y_train_values.float().to(self.device) label = y_train_values.float().to(self.device)
S_mask = S_mask.to(self.device) S_mask = S_mask.to(self.device)
priors = 1-S_mask priors = 1 - S_mask
(vec, sparse_loss) = self.tabnet_model(feature, priors) (vec, sparse_loss) = self.tabnet_model(feature, priors)
f = self.tabnet_decoder(vec) f = self.tabnet_decoder(vec)
@@ -339,9 +360,9 @@ class TabNet_Model(Model):
Pretrain loss function defined in the original paper, read "Tabular self-supervised learning" in https://arxiv.org/pdf/1908.07442.pdf Pretrain loss function defined in the original paper, read "Tabular self-supervised learning" in https://arxiv.org/pdf/1908.07442.pdf
""" """
down_mean = torch.mean(f, dim=0) down_mean = torch.mean(f, dim=0)
down = torch.sqrt(torch.sum(torch.square(f-down_mean), dim = 0)) down = torch.sqrt(torch.sum(torch.square(f - down_mean), dim=0))
up = (f_hat - f)*S up = (f_hat - f) * S
return torch.sum(torch.square(up/down)) return torch.sum(torch.square(up / down))
def loss_fn(self, pred, label): def loss_fn(self, pred, label):
mask = ~torch.isnan(label) mask = ~torch.isnan(label)
@@ -364,12 +385,14 @@ 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
class DecoderStep(nn.Module): class DecoderStep(nn.Module):
@@ -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):
""" """
@@ -395,11 +416,11 @@ class TabNet_Decoder(nn.Module):
super().__init__() super().__init__()
if n_shared > 0: if n_shared > 0:
self.shared = nn.ModuleList() self.shared = nn.ModuleList()
self.shared.append(nn.Linear(inp_dim, 2*out_dim)) self.shared.append(nn.Linear(inp_dim, 2 * out_dim))
for x in range(n_shared - 1): for x in range(n_shared - 1):
self.shared.append(nn.Linear(out_dim, 2*out_dim)) # preset the linear function we will use self.shared.append(nn.Linear(out_dim, 2 * out_dim)) # preset the linear function we will use
else: else:
self.shared=None self.shared = None
self.n_steps = n_steps self.n_steps = n_steps
self.steps = nn.ModuleList() self.steps = nn.ModuleList()
for x in range(n_steps): for x in range(n_steps):
@@ -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
@@ -434,13 +456,13 @@ class TabNet(nn.Module):
self.shared = nn.ModuleList() self.shared = nn.ModuleList()
self.shared.append(nn.Linear(inp_dim, 2 * (n_d + n_a))) self.shared.append(nn.Linear(inp_dim, 2 * (n_d + n_a)))
for x in range(n_shared - 1): for x in range(n_shared - 1):
self.shared.append(nn.Linear(n_d + n_a, 2 * (n_d + n_a))) # preset the linear function we will use self.shared.append(nn.Linear(n_d + n_a, 2 * (n_d + n_a))) # preset the linear function we will use
else: else:
self.shared=None self.shared = None
self.first_step = FeatureTransformer(inp_dim, n_d + n_a, self.shared, n_ind, vbs, device) self.first_step = FeatureTransformer(inp_dim, n_d + n_a, self.shared, n_ind, vbs, device)
self.steps = nn.ModuleList() self.steps = nn.ModuleList()
for x in range(n_steps-1): for x in range(n_steps - 1):
self.steps.append(DecisionStep(inp_dim, n_d, n_a, self.shared, n_ind, relax, vbs, device)) self.steps.append(DecisionStep(inp_dim, n_d, n_a, self.shared, n_ind, relax, vbs, device))
self.fc = nn.Linear(n_d, out_dim) self.fc = nn.Linear(n_d, out_dim)
self.bn = nn.BatchNorm1d(inp_dim, momentum=0.01) self.bn = nn.BatchNorm1d(inp_dim, momentum=0.01)
@@ -449,13 +471,13 @@ class TabNet(nn.Module):
def forward(self, x, priors): def forward(self, x, priors):
assert not torch.isnan(x).any() assert not torch.isnan(x).any()
x = self.bn(x) x = self.bn(x)
x_a = self.first_step(x)[:, self.n_d:] x_a = self.first_step(x)[:, self.n_d :]
sparse_loss = torch.zeros(1).to(x.device) sparse_loss = torch.zeros(1).to(x.device)
out = torch.zeros(x.size(0), self.n_d).to(x.device) out = torch.zeros(x.size(0), self.n_d).to(x.device)
for step in self.steps: for step in self.steps:
x_te, l = step(x, x_a, priors) x_te, l = step(x, x_a, priors)
out += F.relu(x_te[:, :self.n_d]) #split the feautre from feat_transformer out += F.relu(x_te[:, : self.n_d]) # split the feautre from feat_transformer
x_a = x_te[:, self.n_d:] x_a = x_te[:, self.n_d :]
sparse_loss += l sparse_loss += l
return self.fc(out), sparse_loss return self.fc(out), sparse_loss
@@ -468,13 +490,14 @@ 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)
self.vbs = vbs self.vbs = vbs
def forward(self, x): def forward(self, x):
chunk = torch.chunk(x, x.size(0)//self.vbs,0) chunk = torch.chunk(x, x.size(0) // self.vbs, 0)
res = [self.bn(y) for y in chunk] res = [self.bn(y) for y in chunk]
return torch.cat(res, 0) return torch.cat(res, 0)
@@ -486,18 +509,19 @@ 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:
self.fc = fc self.fc = fc
else: else:
self.fc = nn.Linear(inp_dim, out_dim*2) self.fc = nn.Linear(inp_dim, out_dim * 2)
self.bn = GBN(out_dim * 2, vbs=vbs) self.bn = GBN(out_dim * 2, vbs=vbs)
self.od = out_dim self.od = out_dim
def forward(self, x): def forward(self, x):
x = self.bn(self.fc(x)) x = self.bn(self.fc(x))
return torch.mul(x[:, :self.od], torch.sigmoid(x[:, self.od:])) return torch.mul(x[:, : self.od], torch.sigmoid(x[:, self.od :]))
class AttentionTransformer(nn.Module): class AttentionTransformer(nn.Module):
@@ -507,17 +531,18 @@ 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)
self.bn = GBN(inp_dim, vbs=vbs) self.bn = GBN(inp_dim, vbs=vbs)
self.r = relax self.r = relax
#a:feature from previous decision step # a:feature from previous decision step
def forward(self, a, priors): def forward(self, a, priors):
a = self.bn(self.fc(a)) a = self.bn(self.fc(a))
mask = SparsemaxFunction.apply(a * priors) mask = SparsemaxFunction.apply(a * priors)
priors = priors * (self.r - mask) #updating the prior priors = priors * (self.r - mask) # updating the prior
return mask return mask
@@ -528,17 +553,17 @@ class FeatureTransformer(nn.Module):
self.shared = nn.ModuleList() self.shared = nn.ModuleList()
if shared: if shared:
self.shared.append(GLU(inp_dim, out_dim, shared[0], vbs=vbs)) self.shared.append(GLU(inp_dim, out_dim, shared[0], vbs=vbs))
first= False first = False
for fc in shared[1:]: for fc in shared[1:]:
self.shared.append(GLU(out_dim, out_dim, fc, vbs=vbs)) self.shared.append(GLU(out_dim, out_dim, fc, vbs=vbs))
else: else:
self.shared = None self.shared = None
self.independ = nn.ModuleList() self.independ = nn.ModuleList()
if first: if first:
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,22 +577,21 @@ 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)
self.fea_tran = FeatureTransformer(inp_dim, n_d + n_a, shared, n_ind, vbs, device) self.fea_tran = FeatureTransformer(inp_dim, n_d + n_a, shared, n_ind, vbs, device)
def forward(self, x, a, priors): def forward(self, x, a, priors):
mask = self.atten_tran(a, priors) mask = self.atten_tran(a, priors)
sparse_loss = ((-1)*mask*torch.log(mask+1e-10)).mean() sparse_loss = ((-1) * mask * torch.log(mask + 1e-10)).mean()
x = self.fea_tran(x * mask) x = self.fea_tran(x * mask)
return x ,sparse_loss return x, sparse_loss
def make_ix_like(input, dim=0): def make_ix_like(input, dim=0):
@@ -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