| 23 | torch.backends.cudnn.benchmark = False |
| 24 | |
| 25 | def create_simple_model(num_classes=10): |
| 26 | return nn.Sequential( |
| 27 | nn.Conv2d(3, 32, 3, padding=1), |
| 28 | nn.BatchNorm2d(32), |
| 29 | nn.ReLU(inplace=True), |
| 30 | nn.MaxPool2d(2), |
| 31 | |
| 32 | nn.Conv2d(32, 64, 3, padding=1), |
| 33 | nn.BatchNorm2d(64), |
| 34 | nn.ReLU(inplace=True), |
| 35 | nn.MaxPool2d(2), |
| 36 | |
| 37 | nn.Conv2d(64, 128, 3, padding=1), |
| 38 | nn.BatchNorm2d(128), |
| 39 | nn.ReLU(inplace=True), |
| 40 | nn.AdaptiveAvgPool2d(1), |
| 41 | |
| 42 | nn.Flatten(), |
| 43 | nn.Linear(128, num_classes) |
| 44 | ) |
| 45 | |
| 46 | def main(): |
| 47 | parser = argparse.ArgumentParser( |