Upload folder using huggingface_hub

9e55601 verified over 1 year ago

1.91 kB

import torch
import torch.nn as nn

# FT-Transformer Implementation
class FTTransformer(nn.Module):
    def __init__(self, num_features, cat_cardinalities, num_classes, d_token=192, n_blocks=6, attention_dropout=0.2, ffn_dropout=0.2):
        super(FTTransformer, self).__init__()

        self.num_classes = num_classes  # 클래스 개수 저장

        # Embedding layers for categorical features
        self.cat_embeddings = nn.ModuleList([
            nn.Embedding(num_categories, d_token) for num_categories in cat_cardinalities
        ])

        # Linear layer for numerical features
        self.num_linear = nn.Linear(num_features, d_token)

        # Transformer blocks
        self.transformer_blocks = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_token,
                nhead=8,
                dim_feedforward=4 * d_token,
                dropout=attention_dropout,
                activation='gelu'
            ) for _ in range(n_blocks)
        ])

        self.ffn_dropout = nn.Dropout(ffn_dropout)
        if num_classes == 2:
            self.output_layer = nn.Linear(d_token, 1)  # Binary classification
        elif num_classes > 2:
            self.output_layer = nn.Linear(d_token, num_classes)  # Multi classification

    def forward(self, x_num, x_cat):
        # Numerical feature embedding
        x_num = self.num_linear(x_num)

        # Categorical feature embedding
        x_cat = [embed(x_cat[:, i]) for i, embed in enumerate(self.cat_embeddings)]
        x_cat = torch.stack(x_cat, dim=1)

        # Combine numerical and categorical embeddings
        x = x_num.unsqueeze(1) + x_cat

        # Pass through transformer blocks
        for block in self.transformer_blocks:
            x = block(x)

        # Pooling and output
        x = x.mean(dim=1)
        x = self.ffn_dropout(x)
        x = self.output_layer(x)

        return x

	import torch
	import torch.nn as nn

	# FT-Transformer Implementation
	class FTTransformer(nn.Module):
	def __init__(self, num_features, cat_cardinalities, num_classes, d_token=192, n_blocks=6, attention_dropout=0.2, ffn_dropout=0.2):
	super(FTTransformer, self).__init__()

	self.num_classes = num_classes # 클래스 개수 저장

	# Embedding layers for categorical features
	self.cat_embeddings = nn.ModuleList([
	nn.Embedding(num_categories, d_token) for num_categories in cat_cardinalities
	])

	# Linear layer for numerical features
	self.num_linear = nn.Linear(num_features, d_token)

	# Transformer blocks
	self.transformer_blocks = nn.ModuleList([
	nn.TransformerEncoderLayer(
	d_model=d_token,
	nhead=8,
	dim_feedforward=4 * d_token,
	dropout=attention_dropout,
	activation='gelu'
	) for _ in range(n_blocks)
	])

	self.ffn_dropout = nn.Dropout(ffn_dropout)
	if num_classes == 2:
	self.output_layer = nn.Linear(d_token, 1) # Binary classification
	elif num_classes > 2:
	self.output_layer = nn.Linear(d_token, num_classes) # Multi classification

	def forward(self, x_num, x_cat):
	# Numerical feature embedding
	x_num = self.num_linear(x_num)

	# Categorical feature embedding
	x_cat = [embed(x_cat[:, i]) for i, embed in enumerate(self.cat_embeddings)]
	x_cat = torch.stack(x_cat, dim=1)

	# Combine numerical and categorical embeddings
	x = x_num.unsqueeze(1) + x_cat

	# Pass through transformer blocks
	for block in self.transformer_blocks:
	x = block(x)

	# Pooling and output
	x = x.mean(dim=1)
	x = self.ffn_dropout(x)
	x = self.output_layer(x)

	return x