first code commi. First protoyp for AIIA Config
This commit is contained in:
parent
d37124219c
commit
7287ba543f
|
@ -0,0 +1,30 @@
|
|||
import json
|
||||
|
||||
class AIIAConfig:
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "AIIA",
|
||||
radius: int = 3,
|
||||
activation_function: str = "gelu",
|
||||
hidden_size: int = 128,
|
||||
num_hidden_layers: int = 2,
|
||||
num_channels: int = 3,
|
||||
):
|
||||
self.model_name = model_name
|
||||
self.radius = radius
|
||||
self.activation_function = activation_function
|
||||
self.hidden_size = hidden_size
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_channels = num_channels
|
||||
|
||||
def save(self, file_path):
|
||||
# Save config to JSON
|
||||
with open(file_path, 'w') as f:
|
||||
json.dump(self.__dict__, f)
|
||||
|
||||
@classmethod
|
||||
def load(cls, file_path):
|
||||
# Load config from JSON
|
||||
with open(file_path, 'r') as f:
|
||||
config_dict = json.load(f)
|
||||
return cls(**config_dict)
|
Loading…
Reference in New Issue