resnet50

sunbing7 · Jun 12, 2023 · 1bcc84c · 1bcc84c
1 parent 839d205
commit 1bcc84c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src_torch/util.py b/src_torch/util.py
@@ -133,6 +133,12 @@ def get_dataloader(dataset, train=True, ratio=1.0, batch_size=128):
     elif dataset == 'mnistm':
         dataset = CustomMNISTMDataSet(data_root, train, transform)
     elif dataset == 'asl' or dataset == 'caltech':
+        if dataset == 'caltech':
+            transforms_list = []
+            transforms_list.append(transforms.ToTensor())
+            transforms.CenterCrop(size=224),
+            # transforms_list.append(transforms.Normalize(_mean[dataset], _std[dataset]))
+            transform = transforms.Compose(transforms_list)
         dataset = get_custom_folder(data_root + '/' + str(dataset), train, transform)
     else:
         raise Exception('Invalid dataset')