Pytorch-数据保存载入

pytorch中保存模型相关的函数有3个:

torch.save:利用python的pickle模块实现序列化并保存序列化后的object

torch.load:利用pickle将保存的object反序列化

torch.nn.Module.load_state_dict:通过反序列化得到的state_dict读取保存的训练参数

有两种方法保存模型:

1. torch.save(model, path) # 直接保存整个模型
2. torch.save(model.state_dict(), path) # 保存模型的参数

相应地有两种方法加载保存的模型:

model = torch.load(path) # 直接加载模型
model = Model()                         # 先初始化一个模型
model.load_state_dict(torch.load(path)) # 再加载模型参数

  state_dict

 net._modules.items()

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
在 PyTorch 中,可以使用 `torch.utils.data.Dataset` 类来自定义数据集,并使用 `torch.utils.data.DataLoader` 类来加载数据集。 下面是一个简单的示例,假设我们有一个包含图像和标签的数据集,图像存储在 `images` 文件夹中,标签存储在 `labels.txt` 文件中,每行包含一个标签值(整数): ```python import torch from PIL import Image class CustomDataset(torch.utils.data.Dataset): def __init__(self, root_dir, transform=None): self.root_dir = root_dir self.transform = transform self.labels = [] with open('labels.txt', 'r') as f: lines = f.readlines() self.labels = [int(line.strip()) for line in lines] def __len__(self): return len(self.labels) def __getitem__(self, idx): img_path = f'{self.root_dir}/{idx}.jpg' img = Image.open(img_path).convert('RGB') label = self.labels[idx] if self.transform: img = self.transform(img) return img, label # 创建数据集实例 dataset = CustomDataset('images') # 创建数据集加载器 dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True) ``` 在上面的示例中,`CustomDataset` 类定义了自定义数据集的加载方式。`__init__` 方法用于初始化数据集实例,读取标签文件并保存标签。`__len__` 方法返回数据集的大小,即样本数量。`__getitem__` 方法用于获取一个样本,其中 `idx` 表示样本的索引,根据索引读取该样本的图像文件和标签文件,并返回图像和标签。 我们使用 `dataset = CustomDataset('images')` 创建数据集实例,然后使用 `dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)` 创建数据集加载器。其中,`batch_size` 表示每个批次的样本数量,`shuffle` 表示是否随机打乱数据集。 使用 `dataloader` 加载数据集时,会自动调用 `CustomDataset` 中的 `__getitem__` 方法获取样本,然后将多个样本拼接成一个批次返回。因此,可以直接使用 `for batch_x, batch_y in dataloader:` 的方式遍历数据集。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值