Neural Network in TensorFlow

NeuralNetwork.py

import tensorflow as tf

# Define the model using TensorFlow's low-level API
class CustomModel(tf.Module):
    def __init__(self):
        # Initialize weights for each layer
        self.dense1_weights = tf.Variable(tf.random.normal([18, 10]), name='dense1_weights')
        self.dense1_biases = tf.Variable(tf.zeros([10]), name='dense1_biases')

        self.dense2_weights = tf.Variable(tf.random.normal([10, 20]), name='dense2_weights')
        self.dense2_biases = tf.Variable(tf.zeros([20]), name='dense2_biases')

        self.output_weights = tf.Variable(tf.random.normal([20, 1]), name='output_weights')
        self.output_biases = tf.Variable(tf.zeros([1]), name='output_biases')

    def __call__(self, x):
        # First hidden layer
        x = tf.nn.relu(tf.matmul(x, self.dense1_weights) + self.dense1_biases)
        
        # Second hidden layer
        x = tf.nn.relu(tf.matmul(x, self.dense2_weights) + self.dense2_biases)
        
        # Output layer
        x = tf.matmul(x, self.output_weights) + self.output_biases
        
        return x

# Instantiate the model
model2 = CustomModel()

# Define the loss function
def loss_fn(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_true - y_pred))

# Define the optimizer
optimizer = tf.optimizers.Adam()

# Training step
def train_step(x_batch, y_batch):
    with tf.GradientTape() as tape:
        y_pred = model2(x_batch)
        loss = loss_fn(y_batch, y_pred)
    
    gradients = tape.gradient(loss, model2.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model2.trainable_variables))
    
    return loss

# Example usage
# x_train and y_train should be your input features and labels
# x_train = ...
# y_train = ...

# for epoch in range(epochs):
#     loss = train_step(x_train, y_train)
#     print(f"Epoch {epoch}, Loss: {loss.numpy()}")