update example x 3

README.md

@@ -15,7 +15,7 @@ I built EfficientNetV2.5 s to outperform the existing EfficientNet b0 to b4, Eff
   - Params: 16.64 M
   - Multiply-Add Operations: 5.32 G
   - Image size: train = 299x299 / 304x304, test = 304x304
-  - Classification layer: defaults to 1,000 classes, but is empty
+  - Classification layer: defaults to 1,000 classes
 - **Papers:**
   - EfficientNetV2: Smaller Models and Faster Training: https://arxiv.org/abs/2104.00298
   - Layer-adaptive sparsity for the Magnitude-based Pruning: https://arxiv.org/abs/2010.07611
@@ -42,7 +42,7 @@ print_layer_stats = True    # prints the statistics for each layer of the model
 verbose = True              # prints additional info about the MAC calculation
 
 # Download the model. Skip this step if already downloaded
-base_model = "efficientnetv2.5_rw_s_16.64M_5.27G_in3"
+base_model = "efficientnetv2.5_base_in1k.pth"
 url = f"https://huggingface.co/FredZhang7/efficientnetv2.5_rw_s/resolve/main/{model_name}.pth"
 file_name = f"./{base_model}.pth"
 urllib.request.urlretrieve(url, file_name)
@@ -52,8 +52,8 @@ model.classifier = torch.nn.Linear(in_features=1984, out_features=nclass, bias=T
 macs, nparams = get_model_complexity_info(model, input_size, as_strings=False, print_per_layer_stat=print_layer_stats, verbose=verbose)
 traced_model = torch.jit.trace(model, example_inputs)
 
-# model_name = f'{base_model}_{"{:.2f}".format(nparams / 1e6)}M_{"{:.2f}".format(macs / 1e9)}G.pth'
-traced_model.save(base_model)
+model_name = f'{base_model}_{"{:.2f}".format(nparams / 1e6)}M_{"{:.2f}".format(macs / 1e9)}G.pth'
+traced_model.save(model_name)
 
 # Load the trainable model
 model = torch.load(model_name)