Fix SLSA for ML workflows (#299)

* Fix TF Signed-off-by: Mihai Maruseac <[email protected]> * Fix PT Signed-off-by: Mihai Maruseac <[email protected]> * Only run on linux Signed-off-by: Mihai Maruseac <[email protected]> * Fix lint Signed-off-by: Mihai Maruseac <[email protected]> * Properly exclude the other OSes Signed-off-by: Mihai Maruseac <[email protected]> --------- Signed-off-by: Mihai Maruseac <[email protected]>
sigstore · Aug 22, 2024 · e1e36f0 · e1e36f0
1 parent 79b5e91
commit e1e36f0
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 39 deletions.
diff --git a/.github/workflows/slsa_for_ml.yml b/.github/workflows/slsa_for_ml.yml
@@ -47,18 +47,18 @@ jobs:
     strategy:
       fail-fast: false # Don't cancel other jobs if one fails
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest] ## , macos-latest, windows-latest]
         include:
-          - os: macos-latest
-            os_family: Darwin
+          #- os: macos-latest
+          #  os_family: Darwin
           - os: ubuntu-latest
             os_family: Linux
-          - os: windows-latest
-            os_family: Windows
+          #- os: windows-latest
+          #  os_family: Windows
     outputs:
       hash-ubuntu-latest: ${{ steps.hash.outputs.hash-ubuntu-latest }}
-      hash-macos-latest: ${{ steps.hash.outputs.hash-macos-latest }}
-      hash-windows-latest: ${{ steps.hash.outputs.hash-windows-latest }}
+      # hash-macos-latest: ${{ steps.hash.outputs.hash-macos-latest }}
+      # hash-windows-latest: ${{ steps.hash.outputs.hash-windows-latest }}
     steps:
     - run: git config --global core.autocrlf input
     - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
@@ -101,7 +101,7 @@ jobs:
     strategy:
       fail-fast: false # Don't cancel other jobs if one fails
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest] #, macos-latest, windows-latest]
     permissions:
       actions: read
       id-token: write

diff --git a/slsa_for_models/pytorch_cifar10.py b/slsa_for_models/pytorch_cifar10.py
@@ -12,32 +12,55 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
+import torch.nn as nn
 
-# We will do a lazy import for these 7 modules, exploiting Python's symbol
+
+# We will do a lazy import for these 5 modules, exploiting Python's symbol
 # resolution. The lazy import is needed to make sure we only import PyTorch
 # libraries only if we want to train a PyTorch model.
-torch = None
-nn = None
 functional = None
 optim = None
 torchvision = None
 transforms = None
 
 
+class MyModel(nn.Module):
+    """Train a PyTorch model.
+
+    Based on tutorial from
+    https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = self.pool(functional.relu(self.conv1(x)))
+        x = self.pool(functional.relu(self.conv2(x)))
+        x = torch.flatten(x, 1)
+        x = functional.relu(self.fc1(x))
+        x = functional.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
 def pretraining():
     """Perform setup required before training.
 
     Does the lazy loading of TensorFlow too, to prevent compatibility issues
     with mixing TensorFlow and PyTorch imports.
     """
-    global torch
-    global nn
     global functional
     global optim
     global torchvision
     global transforms
-    import torch
-    import torch.nn as nn
     import torch.nn.functional as functional
     import torch.optim as optim
     import torchvision
@@ -89,29 +112,6 @@ def create_model():
 
     Returns the model.
     """
-
-    # Train a model based on tutorial from
-    # https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html.
-    # We inline the class to be able to use lazy loading of PyTorch modules.
-    class MyModel(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.conv1 = nn.Conv2d(3, 6, 5)
-            self.pool = nn.MaxPool2d(2, 2)
-            self.conv2 = nn.Conv2d(6, 16, 5)
-            self.fc1 = nn.Linear(16 * 5 * 5, 120)
-            self.fc2 = nn.Linear(120, 84)
-            self.fc3 = nn.Linear(84, 10)
-
-        def forward(self, x):
-            x = self.pool(functional.relu(self.conv1(x)))
-            x = self.pool(functional.relu(self.conv2(x)))
-            x = torch.flatten(x, 1)
-            x = functional.relu(self.fc1(x))
-            x = functional.relu(self.fc2(x))
-            x = self.fc3(x)
-            return x
-
     return MyModel()
 
 

diff --git a/slsa_for_models/tensorflow_cifar10.py b/slsa_for_models/tensorflow_cifar10.py
@@ -115,7 +115,13 @@ def train_model(model, train, test):
     """
     x, y = train
     model.fit(
-        x, y, batch_size=256, epochs=16, validation_data=test, shuffle=True
+        x,
+        y,
+        batch_size=256,
+        epochs=16,
+        validation_data=test,
+        shuffle=True,
+        verbose=0,
     )