Views
No views yet
1from transformers import AutoModel
2model = AutoModel.from_pretrained("FredZhang7/efficientnetv2.5_rw_s", trust_remote_code=True)pip install ptflops timm1from ptflops import get_model_complexity_info
2import torch
3import urllib.request
4
5nclass = 3 # number of classes in your dataset
6input_size = (3, 304, 304) # recommended image input size
7print_layer_stats = True # prints the statistics for each layer of the model
8verbose = True # prints additional info about the MAC calculation
9
10# Download the model. Skip this step if already downloaded
11base_model = "efficientnetv2.5_base_in1k"
12url = f"https://huggingface.co/FredZhang7/efficientnetv2.5_rw_s/resolve/main/{base_model}.pth"
13file_name = f"./{base_model}.pth"
14urllib.request.urlretrieve(url, file_name)
15
16shape = (2,) + input_size
17example_inputs = torch.randn(shape)
18example_inputs = (example_inputs - example_inputs.min()) / (example_inputs.max() - example_inputs.min())
19
20model = torch.load(file_name)
21model.classifier = torch.nn.Linear(in_features=1984, out_features=nclass, bias=True)
22macs, nparams = get_model_complexity_info(model, input_size, as_strings=False, print_per_layer_stat=print_layer_stats, verbose=verbose)
23traced_model = torch.jit.trace(model, example_inputs)
24
25model_name = f'{base_model}_{"{:.2f}".format(nparams / 1e6)}M_{"{:.2f}".format(macs / 1e9)}G.pth'
26traced_model.save(model_name)
27
28# Load the trainable model
29model = torch.load(model_name)efficientnet_b3_pruned achieved the second highest top-1 accuracy as well as the highest epoch-1 training accuracy on my task, out of EfficientNetV2.5 small and all existing EfficientNet models my 24 GB VRAM RTX 3090 could handle.