智能数字图像处理：VGGNet代码（pytorch）之train.py解读

最新推荐文章于 2022-11-05 20:44:17 发布

王壹浪

最新推荐文章于 2022-11-05 20:44:17 发布

阅读量1k

点赞数

分类专栏：心得人工智能文章标签：深度学习 python 定位 tensorflow

本文链接：https://blog.csdn.net/com_fang_bean/article/details/108491602

版权

本文介绍了VGGNet训练过程中的关键步骤，包括如何实例化模型，如vgg11、13、16、19四种结构，以及使用nn.CrossEntropyLoss作为损失函数和Adam优化器进行训练。

摘要由CSDN通过智能技术生成

唯一与其他模型的train.py不同的是VGGNet有一个实例化模型的方法，即能选择model.py中不同的模型进行实例化再训练。

model_name = "vgg16"
net = vgg(model_name=model_name, num_classes=5, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)

根据上一次博客所讲的这里可以选择vgg11、13、16、19四种网络模型。

import torch.nn as nn
from torchvision import transforms, datasets
import json
import os
import torch.optim as optim
from model import vgg
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

data_transform = {
    "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),