Molly/enable unet1d

torchbenchmark/models/unet1d/__init__.py

@@ -0,0 +1,27 @@
+from ...util.model import BenchmarkModel
+
+import torch
+from diffusers import UNet1DModel
+
+class Model(BenchmarkModel):
+    DEFAULT_EVAL_BSIZE = 32
+    def __init__(self, test, device, batch_size=None, extra_args=[]):
+        super().__init__( test=test, device=device,
+                batch_size=batch_size, extra_args=extra_args)
+
+        self.in_channels = 32
+        self.seq_len = 256
+        self.num_features = 16
+        self.block_out_channels = (self.seq_len, 64, 64)
+        print(self.batch_size)
+        self.example_inputs =  torch.randn(1, self.num_features, self.seq_len)
+        self.timesteps = torch.tensor([1])
+        self.model = UNet1DModel(in_channels=self.in_channels, block_out_channels=self.block_out_channels)
+
+    def get_module(self):
+        return self.model, self.example_inputs
+
+    def eval(self):
+        self.model.eval()
+        with torch.no_grad():
+            out=self.model(self.example_inputs, self.timesteps)

torchbenchmark/models/unet1d/install.py

@@ -0,0 +1,9 @@
+import subprocess
+import sys
+import os
+
+def pip_install_requirements():
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-q', '-r', 'requirements.txt'])
+
+if __name__ == '__main__':
+    pip_install_requirements()

torchbenchmark/models/unet1d/metadata.yml

@@ -0,0 +1,8 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 64
+eval_benchmark: true
+eval_deterministic: false
+eval_nograd: true
+train_benchmark: false
+train_deterministic: false

torchbenchmark/models/unet1d/requirements.txt

@@ -0,0 +1 @@
+diffusers