Pytorch-读数据 or 模型

最新推荐文章于 2024-04-28 17:24:01 发布

原创最新推荐文章于 2024-04-28 17:24:01 发布 · 284 阅读

0 ·

CC 4.0 BY-SA版权

Pytorch 专栏收录该内容

1 篇文章

订阅专栏

本文详细介绍了如何使用PyTorch进行模型训练，包括定义损失函数、优化器、数据加载器，以及如何保存训练好的模型。通过示例展示了AlexNet模型的训练流程，涵盖了模型初始化、参数加载、权重信息和网络参数的处理。

val_tf = pretrainedmodels.utils.TransformImage(
    model,
    scale=scale,
    preserve_aspect_ratio=args.preserve_aspect_ratio
)

val_loader = torch.utils.data.DataLoader(
    datasets.ImageFolder(valdir, val_tf),
    batch_size=args.batch_size, shuffle=False,
    num_workers=args.workers, pin_memory=True)

# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()

optimizer = torch.optim.SGD(model.parameters(), args.lr,
                            momentum=args.momentum,
                            weight_decay=args.weight_decay)

model = torch.nn.DataParallel(model).cuda()

https://github.com/Cadene/pretrained-models.pytorch

https://www.e-learn.cn/content/qita/850153

# 保存模型

torch.save(model.state_dict(), alexnet.pth)

# 定义模型

model = AlexNet()

# 根据模型读取参数

model.load_state_dict(state_dict)

# 权重信息

restore_param = {v for k, v in model.state_dict().items() }

# 网络信息