ocean algorithm to do actual inference

Author: devinaconley

mlweb3/ocean/algorithm.py

@@ -6,24 +6,88 @@
 import argparse
 import json
 
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
 
 def main():
     args = parse_arguments()
     print(args)
-
     print(os.getcwd())
 
-    dids = os.getenv('DIDS', None)
-    print(dids)
+    if args['local']:
+        filepath = args['weights']
+    else:
+        dids = os.getenv('DIDS', None)
+        print(dids)
+        if not dids:
+            print('no DIDs found in environment. exiting.')
+            return
+        dids = json.loads(dids)
+        if len(dids) == 0:
+            print('no DID found for model. exiting.')
+        filepath = f'data/inputs/{dids[0]}/0'
+
+    # load model weights
+    print(f'Loading SimpleCNN from {filepath}...')
+    model = SimpleCNN()
+    model.load_state_dict(torch.load(filepath))
+    model.eval()
+
+    # get/load data
+    # TODO use local data when multiple input assets are supported
+    os.makedirs('./etc/mnist', exist_ok=True)
+    data = DataLoader(
+        datasets.MNIST('./etc/mnist', train=False, download=True, transform=ToTensor()),
+        batch_size=64
+    )
+
+    # do inference
+    correct, total = 0, 0
+    predictions = []
+    with torch.no_grad():
+        for X, y in data:
+            # X, y = X.to(device), y.to(device)
+            pred = model(X)
+            correct += (pred.argmax(1) == y).sum().item()
+            total += len(X)
+            predictions.extend(pred.argmax(1).numpy().tolist())
+
+    print(f'test:\n  accuracy: {correct / total:>0.4f}')
+
+    # write output
+    output_file = 'results.txt' if args['local'] else '/data/outputs/result'
+    with open(output_file, 'w') as f:
+        f.write(f'accuracy: {correct / total:>0.4f}\n\npredictions:\n')
+        for p in predictions:
+            f.write(f'{p}\n')
+
 
-    if not dids:
-        print('No DIDs found in environment. Aborting.')
-        return
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 32, (5, 5))
+        self.conv2 = nn.Conv2d(32, 64, (5, 5))
+        self.fc1 = nn.Linear(1024, 128)
+        self.fc2 = nn.Linear(128, 10)
 
-    dids = json.loads(dids)
-    for did in dids:
-        filename = f'data/inputs/{did}/0'
-        print(f'Reading asset file {filename}.')
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.fc2(x)
+        return x
 
 
 def parse_arguments() -> dict:

mlweb3/ocean/inference.py

@@ -78,14 +78,16 @@ def predict():
     assert datasets, 'payment for dataset unsuccessful'
     assert algorithm, 'payment for algorithm unsuccessful'
 
+    # NOTE: not currently possible to use multiple input datasets on a single compute job
+
     # start compute job
     t0 = datetime.datetime.utcnow()
     job_id = ocean.compute.start(
         consumer_wallet=bob_wallet,
-        dataset=data_compute_input,
+        dataset=datasets[0],
         compute_environment=free_c2d_env['id'],
-        algorithm=algo_compute_input,
-        additional_datasets=[weights_compute_input]
+        algorithm=algorithm,
+        # additional_datasets=[datasets[1]]
     )
     print('Started compute job: {}'.format(job_id))