File size: 2,633 Bytes
91f863a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import copy
from collections import OrderedDict
import torch.nn as nn
ACT2CLS = {
"relu": nn.ReLU,
"relu6": nn.ReLU6,
"gelu": nn.GELU,
"elu": nn.ELU,
"sigmoid": nn.Sigmoid,
"tanh": nn.Tanh,
}
class ClassInstantier(OrderedDict):
def __getitem__(self, key):
content = super().__getitem__(key)
cls, kwargs = content if isinstance(content, tuple) else (content, {})
return cls(**kwargs)
ACT2FN = ClassInstantier(ACT2CLS)
def get_actfunction(activation_string):
if activation_string in ACT2FN:
return ACT2FN[activation_string]
else:
raise KeyError(f"function {activation_string} not found in ACT2FN mapping {list(ACT2FN.keys())}")
def deleteEncodingLayers(model, num_layers_to_keep): # must pass in the full bert model
oldModuleList = model.encoder.layer
newModuleList = nn.ModuleList()
# Now iterate over all layers, only keepign only the relevant layers.
for i in range(num_layers_to_keep):
newModuleList.append(oldModuleList[i])
# create a copy of the model, modify it with the new list, and return
copyOfModel = copy.deepcopy(model)
copyOfModel.encoder.layer = newModuleList
return copyOfModel
def make_sequential(input_dim:int, output_dim:int, act_function:str = "relu", dropout:float = 0.1, shift:float = 0.0):
act_func = get_actfunction(act_function.lower())
sequencial = nn.Sequential(nn.Linear(input_dim, output_dim),
act_func,
nn.Dropout(dropout))
return sequencial
def make_encoder(input_dim:int, num_heads:int, num_layers:int, act_function:str = "relu", dropout:float = 0.1):
act_func = get_actfunction(act_function.lower())
# sequence_pos_encoding = PositionalEncoding(input_dim, dropout)
seq_trans_encoder_layer = nn.TransformerEncoderLayer(d_model=input_dim, nhead=num_heads, dim_feedforward=input_dim * num_heads, dropout=dropout, activation=act_func)
seqTransEncoder = nn.TransformerEncoder(seq_trans_encoder_layer, num_layers=num_layers)
return seqTransEncoder
def make_decoder(input_dim:int, num_heads:int, num_layers:int, act_function:str = "relu", dropout:float = 0.1):
act_func = get_actfunction(act_function.lower())
# sequence_pos_encoding = PositionalEncoding(input_dim, dropout)
seq_trans_decoder_layer = nn.TransformerDecoderLayer(d_model=input_dim, nhead=num_heads, dim_feedforward=input_dim * num_heads, dropout=dropout, activation=act_func)
seqTransDecoder = nn.TransformerDecoder(seq_trans_decoder_layer, num_layers=num_layers)
return seqTransDecoder |