| import torch.nn as nn | |
| import timm | |
| def Model(): | |
| model = timm.create_model("mobilenetv3_large_100", pretrained=True) | |
| model.classifier = nn.Linear(1280, 10) | |
| for name, param in model.named_parameters(): | |
| if "bn" in name: | |
| # print(f"freeze {name}") | |
| param.requires_grad = False | |
| return model, model.classifier | |
| if __name__ == "__main__": | |
| model, _ = Model() | |
| print(model) | |
| num_param = 0 | |
| for v in model.parameters(): | |
| num_param += v.numel() | |
| print("num_param:", num_param) | |