Refactoring README to include run-time information (#2275)

images/pytorch-cuda12/README.md

@@ -31,7 +31,7 @@ docker pull cgr.dev/chainguard/pytorch-cuda12:latest
 ## Running pytorch-cuda12
 
 Pytorch has some pre-requisites which need to be configured in the environment
-prior to running. For examples, please refer to [TESTING.md](https://github.com/chainguard-images/images/blob/main/images/pytorch/TESTING.md).
+prior to running with GPUs. For examples, please refer to [TESTING.md](https://github.com/chainguard-images/images/blob/main/images/pytorch-cuda12/TESTING.md).
 
 Additionally, please refer to the [upstream documentation](https://github.com/pytorch/pytorch)
 for more information on coniguring and using Pytorch.
@@ -45,4 +45,36 @@ docker run --rm -i -t \
     --gpus all \
     cgr.dev/chainguard/pytorch:latest
 ```
+
+If your environment has access to GPUs, you may provide access pytorch access to it by running 
+```bash
+docker run --rm -it --gpus all cgr.dev/chainguard/pytorch-cuda12:latest
+bash-5.2$ python
+Python 3.11.8 (main, Feb  7 2024, 00:46:15) [GCC 13.2.0] on linux
+Type "help", "copyright", "credits" or "license" for more information.
+>>> import torch
+>>> print(torch.cuda.is_available())
+True
+```
+As a quick intro, we will use pytorch to create a very simple deep learning model with two linear layers and an activation function. We’ll create an instance of it and ask it to report on its parameters. The script can be found in ```model_builder.py``` in this directory.
+
+To run this script, 
+```bash
+docker run --rm -it -v /home/srishihegde/quick.py:/tmp/model_builder.py --gpus all cgr.dev/chainguard/pytorch-cuda12:latest python /tmp/model_builder.py
+```
+
+### Using Helm charts
+
+As a place to get started, you may also use this Helm chart to get pytorch running
+```bash
+  helm install pytorch \
+  --namespace pytorch-space --create-namespace  \
+  --set image.registry="cgr.dev" \
+  --set image.repository="chainguard/pytorch-cuda12" \
+  --set image.tag=latest \
+  --set containerSecurityContext.runAsUser=0 \
+  --set containerSecurityContext.runAsNonRoot=false \
+  --set containerSecurityContext.allowPrivilegeEscalation=true \
+  --wait oci://registry-1.docker.io/bitnamicharts/pytorch
+```
 <!--body:end-->

images/pytorch-cuda12/config/main.tf

@@ -40,7 +40,7 @@ output "config" {
       "PATH" : "/usr/share/pytorch/.venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
     }, var.environment)
     entrypoint = {
-      command = ""
+      command = "/bin/bash -c"
     }
     archs = ["x86_64"]
   })

images/pytorch-cuda12/model_builder.py

@@ -0,0 +1,34 @@
+import torch
+
+class TinyModel(torch.nn.Module):
+
+    def __init__(self):
+        super(TinyModel, self).__init__()
+
+        self.linear1 = torch.nn.Linear(100, 200)
+        self.activation = torch.nn.ReLU()
+        self.linear2 = torch.nn.Linear(200, 10)
+        self.softmax = torch.nn.Softmax()
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.activation(x)
+        x = self.linear2(x)
+        x = self.softmax(x)
+        return x
+
+tinymodel = TinyModel()
+
+print('The model:')
+print(tinymodel)
+
+print('\n\nJust one layer:')
+print(tinymodel.linear2)
+
+print('\n\nModel params:')
+for param in tinymodel.parameters():
+    print(param)
+
+print('\n\nLayer params:')
+for param in tinymodel.linear2.parameters():
+    print(param)