microsoft / knossos-ksc

Compiler with automatic differentiation
Other
45 stars 10 forks source link

Doc: Make the vrelu example from the intro work. #1025

Open awf opened 2 years ago

awf commented 2 years ago

The doc intro describes this NN - we should make it a proper example, and unit test/benchmark it

We'll make a simple DNN with vrelu3 activations, as in link

# Initialize the model using nn.Sequential
model = nn.Sequential(OrderedDict([
                    ('fc1', nn.Linear(784, 256)),
                    ('activation1', vrelu3),
                    ('fc2', nn.Linear(256, 128)),
                    ('bn2', nn.BatchNorm1d(num_features=128)),
                    ('activation2', vrelu3),
                    ('dropout', nn.Dropout(0.3)),
                    ('fc3', nn.Linear(128, 64)),
                    ('bn3', nn.BatchNorm1d(num_features=64)),
                    ('activation3', vrelu3),
                    ('logits', nn.Linear(64, 10)),
                    ('logsoftmax', nn.LogSoftmax(dim=1))]))

# Run training
train_model(model)