torchvision:https://pytorch.org/vision/stable/models.html
知识点
mymodel.train() /eval() 声明
- 只对Dropout和BatchNorm有作用
现有模型使用和修改
修改层
vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True) # 已经训练好的权重
vgg16_true.classifier.add_module("add_linear", nn.Linear(1000, 10)) # 添加新的层
vgg16_false.classifier[6] = nn.Linear(in_features=4096, out_features=10) # 修改新的层=======================================
添加新的层add_module
vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_true = torchvision.models.vgg16(pretrained=True) # 已经训练好的权重
vgg16_true.classifier.add_module("add_linear", nn.Linear(1000, 10)) # 添加新的层
完整模型训练 套路
# 1 准备数据集
train_data = torchvision.datasets.CIFAR10("../../data", train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10("../../data", train=False, transform=torchvision.transforms.ToTensor(),
download=True)
print("训练集数据长度:%d \n测试集数据长度:%d"%(len(train_data), len(test_data)))
# 训练集数据长度:50000
# 测试集数据长度:10000
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 2 model
mymodule = Mymodule()
# 3 损失函数和优化器
loss_fn = nn.CrossEntropyLoss()
optimizar = torch.optim.SGD(mymodule.parameters(), lr=1e-2)
# 4 训练 circle
## 记录训练的次数
total_train_step = 0
## 记录测试的次数
total_test_step = 0
## 训练论数
epoch = 10
for i in range(epoch):
print("------训练第%d轮"%(i+1))
for data in train_dataloader:
imgs, target = data
output = mymodule(imgs)
loss = loss_fn(output, target)
# 优化器
optimizar.zero_grad()
loss.backward()
optimizar.step()
total_train_step +=1
print("训练次数:%d, Loss:%f"%(total_train_step, loss.item()))