导入
from torchvision import models
#固定参数 不训练
feature_extract = True
def set_parameter_requires_grad(model, feature_extracting): if feature_extracting: for param in model.parameters(): param.requires_grad = False
def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True): model_ft = models.resnet18(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 102) # 类别数自己根据自己任务来 input_size = 64 # 输入大小根据自己配置来 return model_ft, input_size
model_ft, input_size = initialize_model(model_name, 102, feature_extract, use_pretrained=True) model_ft = model_ft.to(device) filename='checkpoint.pth'
params_to_update = model_ft.parameters() print("Params to learn:") if feature_extract: params_to_update = [] for name,param in model_ft.named_parameters(): if param.requires_grad == True: params_to_update.append(param) print("\t",name)
# 优化器设置 optimizer_ft = optim.Adam(params_to_update, lr=1e-2)#要训练啥参数,你来定 scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)#学习率每7个epoch衰减成原来的1/10 criterion = nn.CrossEntropyLoss()