torchvision：https://pytorch.org/vision/stable/models.html

知识点

mymodel.train() /eval() 声明

  • 只对Dropout和BatchNorm有作用

image.png
image.png

现有模型使用和修改

P25现有网络模型的使用及修改

修改层

  1. vgg16_false = torchvision.models.vgg16(pretrained=False)
  2. vgg16_true = torchvision.models.vgg16(pretrained=True) # 已经训练好的权重
  4. vgg16_true.classifier.add_module("add_linear", nn.Linear(1000, 10)) # 添加新的层
  5. vgg16_false.classifier[6] = nn.Linear(in_features=4096, out_features=10)  # 修改新的层=======================================

image.png

添加新的层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)) # 添加新的层

image.png

完整模型训练 套路


# 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()))