test.py

from hotgrad.sequential import Sequential

from hotgrad.variable import Variable
from hotgrad.sequential import Sequential
from hotgrad.functions.layers import Linear
from hotgrad.functions.activations import ReLU, Tanh
from hotgrad.functions.losses import MSE
from hotgrad.optimizers import SGD

from torch import manual_seed

from dataset_generator import generate_dataset

# the seed ensures that the test case is deterministic
manual_seed(3)

# generate the dataset
X_train, X_test, y_train, y_test = generate_dataset(1000, one_hot_encoding=True)
y_train.data = y_train.data*2-1
y_test.data = y_test.data*2-1

X_train_baselines = X_train.data.numpy()
X_test_baselines = X_test.data.numpy()

X_train.data = (X_train.data - X_train.data.mean())/X_train.data.std() 
X_test.data = (X_test.data - X_test.data.mean())/X_test.data.std() 

# model: two input units, two output units, three hidden layers of 25 units
model = Sequential([Linear(25), ReLU(), Linear(25), ReLU(), Linear(25), ReLU(), Linear(2), Tanh()], MSE(), SGD(lr=0.01))
model.fit(X_train, y_train, X_test, y_test, epochs=300, log_error=True, verbose=True)