From fc6597a6d9f3fac6a533fa4f0fd08f52dc8dc6d6 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sat, 30 Oct 2021 16:48:39 -0700 Subject: [PATCH] only resnet18, it's too slow otherwise --- {docs => accel/cherry}/nvidia_notes | 0 test/test_train.py | 4 ++-- 2 files changed, 2 insertions(+), 2 deletions(-) rename {docs => accel/cherry}/nvidia_notes (100%) diff --git a/docs/nvidia_notes b/accel/cherry/nvidia_notes similarity index 100% rename from docs/nvidia_notes rename to accel/cherry/nvidia_notes diff --git a/test/test_train.py b/test/test_train.py index f7edcc17ef..0ccad33cf8 100644 --- a/test/test_train.py +++ b/test/test_train.py @@ -8,7 +8,7 @@ from extra.training import train from extra.utils import get_parameters from models.efficientnet import EfficientNet from models.transformer import Transformer -from models.resnet import ResNet18, ResNet34, ResNet50 +from models.resnet import ResNet18 BS = int(os.getenv("BS", "4")) @@ -41,7 +41,7 @@ class TestTrain(unittest.TestCase): def test_resnet(self): X = np.zeros((BS, 3, 224, 224), dtype=np.float32) Y = np.zeros((BS), dtype=np.int32) - for resnet_v in [ResNet18, ResNet34, ResNet50]: + for resnet_v in [ResNet18]: model = resnet_v(num_classes=1000, pretrained=True) train_one_step(model, X, Y)