Beispiel #1
0
#         self.last_fc = nn.Linear(in_features=in_size, out_features=output_size)

#     def forward(self, x):
#         for i, fc in enumerate(self.fcs):
#             x = fc(x)
#             x = nn.ReLU()(x)
#         out = self.last_fc(x)
#         return nn.Sigmoid()(out)

# Set device to be used
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device used:', device)
model_pytorch = SimpleModel(input_size=input_size,
                            hidden_sizes=hidden_sizes,
                            output_size=output_size)
model_pytorch = model_pytorch.to(device)

# Set loss and optimizer
# Set binary cross entropy loss since 2 classes only
criterion = nn.BCELoss()
optimizer = optim.Adam(model_pytorch.parameters(), lr=1e-3)

num_epochs = 20

# Train model
time_start = time.time()

for epoch in range(num_epochs):
    model_pytorch.train()

    train_loss_total = 0
Beispiel #2
0
import numpy as np
import torch

test_size = 2000

input_size = 20
hidden_sizes = [50, 50]
output_size = 1

X_test = np.random.randn(test_size, input_size).astype(np.float32)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model_pytorch = SimpleModel(input_size=input_size,
                            hidden_sizes=hidden_sizes,
                            output_size=output_size)
model_pytorch = model_pytorch.to(device)  # important...

model_pytorch.load_state_dict(torch.load('./models/model_simple.pt'))

dummy_input = torch.from_numpy(X_test[0].reshape(1, -1)).float().to(device)
dummy_output = model_pytorch(dummy_input)
print(dummy_output)

# Export to ONNX format
torch.onnx.export(model_pytorch,
                  dummy_input,
                  './models/model_simple.onnx',
                  input_names=['test_input'],
                  output_names=['test_output'])