File size: 537 Bytes
f7009b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
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)
|