File size: 1,519 Bytes
16456f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import os
import torch
import torch.nn as nn
import torch.nn.init as init
from .vit import ViT
class ViTLargeClassifier(nn.Module):
def __init__(self, num_classes: int = 14, image_size=224, patch_size=16):
super(ViTLargeClassifier, self).__init__()
# 初始化ViT模型
self.vit = ViT(
image_size=image_size,
patch_size=patch_size,
num_classes=num_classes,
dim=1024,
depth=24,
heads=16,
mlp_dim=4096,
dropout=0.1,
emb_dropout=0.1
)
# 初始化权重
if not self.load():
for m in self.modules():
if isinstance(m, nn.Linear):
init.xavier_normal_(m.weight)
if m.bias is not None:
init.zeros_(m.bias)
elif isinstance(m, nn.LayerNorm):
init.ones_(m.weight)
init.zeros_(m.bias)
def forward(self, x):
return self.vit(x)
def load(self, filename: str = None) -> bool:
if filename is None:
current_work_dir = os.path.dirname(__file__)
filename = os.path.join(current_work_dir, "best_pth", "ViTLargeClassifier.pth")
if not os.path.exists(filename):
print("Model file does not exist.")
return False
self.load_state_dict(torch.load(filename))
print("Model loaded successfully.")
return True
|