一个通用的pytorch使用预训练模型训练本地数据代码--resnet101为例

最新推荐文章于 2024-10-30 16:54:53 发布

寂寞军刀

最新推荐文章于 2024-10-30 16:54:53 发布

阅读量1.5k

点赞数 3

分类专栏：深度学习文章标签： pytorch 神经网络深度学习

本文链接：https://blog.csdn.net/akansase/article/details/111305315

版权

本文介绍了如何利用PyTorch结合预训练的ResNet101模型，对本地数据进行深度学习训练。代码中包含详细注释，适用于PyTorch1.6及以上版本，对于更低版本需关注除法计算差异。

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

因为代码里面有比较详细的注记，这里就不啰嗦了。注意，如果是pytorch1.6以下版本，在计算除法时会有不同，具体查看该版本下除法计算的文档。直接上干货了。

from torch import  nn
from torchvision import models,datasets,transforms
import torch as tc
from  torch.autograd import Variable
from torch.utils.data import DataLoader
from matplotlib import pyplot as plt
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
batch_size = 8
learning_rate = 0.0002
epoc = 10

train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])
val_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])

train_dir = 'E:/pydata/cats_and_dogs_small/train'
train_datasets = datasets.ImageFolder(train_dir,transform=train_transforms)
train_dataloader = DataLoader(train_datasets,batch_size=batch_size,shuffle=True)

val_dir