File size: 384 Bytes
f7009b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import torch
import torch.nn as nn
import timm
def Model():
model = timm.create_model("vit_tiny_patch16_224", pretrained=False)
model.head = nn.Linear(192, 10)
return model, model.head
if __name__ == "__main__":
model, _ = Model()
print(model)
num_param = 0
for v in model.parameters():
num_param += v.numel()
print("num_param:", num_param) |