基于卷积神经网络的猫狗识别

TensorFlow和Keras

数据来源：Kaggle在2013年公开的猫狗数据集，该数据集总共25000张图片，猫狗各12500张。
下载链接:https://www.kaggle.com/c/dogs-vs-cats/data
代码：

import os,shutil

original_dataset_diar = '/home/u/notebook_workspase/datas/dogs-vs-cats/train'#原始数据解压目录

base_dir = '/home/u/notebook_workspase/datas/dogs-cats-small-dataset'#自己保留的小数据集
os.mkdir(base_dir)

#划分后的train，validation，test目录
train_dir = os.path.join(base_dir,'train')#将多个路径组合后返回
os.mkdir(train_dir)
validation_dir = os.path.join(base_dir,'validation')
os.mkdir(validation_dir)
test_dir = os.path.join(base_dir,'test')
os.mkdir(test_dir)
#猫和狗的train，validation，test图像目录
train_cats_dir = os.path.join(train_dir,'cats')
os.mkdir(train_cats_dir)
train_dogs_dir = os.path.join(train_dir,'dogs')
os.mkdir(train_dogs_dir)

validation_cats_dir = os.path.join(validation_dir,'cats')
os.mkdir(validation_cats_dir)
validation_dogs_dir = os.path.join(validation_dir,'dogs')
os.mkdir(validation_dogs_dir)

test_cats_dir = os.path.join(test_dir,'cats')
os.mkdir(test_cats_dir)
test_dogs_dir = os.path.join(test_dir,'dogs')
os.mkdir(test_dogs_dir)
# 复制1000猫到训练目录中
fnames = ['cat.{}.jpg'.format(i) for i in range(1000)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(train_cats_dir,fname)
    shutil.copyfile(src,dst)
# 500张猫的验证图片，依次类推
fnames = ['cat.{}.jpg'.format(i) for i in range(1000,1500)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(validation_cats_dir,fname)
    shutil.copyfile(src,dst)
fnames = ['cat.{}.jpg'.format(i) for i in range(1500,2000)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(test_cats_dir,fname)
    shutil.copyfile(src,dst)
fnames = ['dog.{}.jpg'.format(i) for i in range(1000)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(train_dogs_dir,fname)
    shutil.copyfile(src,dst)
fnames = ['dog.{}.jpg'.format(i) for i in range(1000,1500)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(validation_dogs_dir,fname)
    shutil.copyfile(src,dst)
fnames = ['dog.{}.jpg'.format(i) for i in range(1500,2000)]
for fname in fnames:
    src = os.path.join(original_dataset_diar,fname)
    dst = os.path.join(test_dogs_dir,fname)
    shutil.copyfile(src,dst)
print(len(os.listdir(train_cats_dir)))
print(len(os.listdir(train_dogs_dir)))
print(len(os.listdir(validation_cats_dir)))
print(len(os.listdir(validation_dogs_dir)))
print(len(os.listdir(test_dogs_dir)))
from keras import layers
from keras import models

model = models.Sequential()
model.add(layers.Conv2D(32,(3,3),activation='relu',input_shape=(150,150,3)))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(64,(3,3),activation='relu'))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(128,(3,3),activation='relu',))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(128,(3,3),activation='relu',))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Flatten())
model.add(layers.Dense(512,activation='relu'))
model.add(layers.Dense(1,activation='sigmoid'))
Using TensorFlow backend.
model.summary()
# 编译
from keras import optimizers

model.compile(loss='binary_crossentropy',
             optimizer=optimizers.RMSprop(lr = 1e-4),
             metrics=['acc'])
from keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
                train_dir, # 目标目录
                target_size=(150, 150), # 所有图像调整为150x150
                batch_size=20,
                class_mode='binary') # 二进制标签
validation_generator = test_datagen.flow_from_directory(
                validation_dir,
                target_size=(150, 150), 
                batch_size=20,
                class_mode='binary')
history = model.fit_generator(
    train_generator,#python 生成器
    steps_per_epoch=100,#100批次
    epochs=30,
    validation_data=validation_generator,
    validation_steps=50)
model.save('cat-dog-small-1.h5')#保存模型
import matplotlib.pyplot as plt
%matplotlib inline

loss = history.history['loss']
val_loss = history.history['val_loss']
acc = history.history['acc']
val_acc = history.history['val_acc']

epochs = range(1,len(acc)+1)

plt.plot(epochs,acc,'bo',label = 'Training acc')
plt.plot(epochs,val_acc,'b',label = 'Validation acc')
plt.title('Training and validation accuracy')
plt.legend()#显示标签

plt.figure()

plt.plot(epochs,loss,'bo',label = 'Training loss')
plt.plot(epochs,val_loss,'b',label = 'Validation loss')
plt.title("Training and validation loss")
plt.legend()

plt.show()

datagen = ImageDataGenerator(
    rotation_range=40, # 图像随机旋转的角度范围
    width_shift_range=0.2, # 水平或垂直平移范围，相对总宽度或总高度的比例的比例
    height_shift_range=0.2, 
    shear_range=0.2, # 随机错切变换角度换角度
    zoom_range=0.2, # 随机缩放范围
    horizontal_flip=True, # 一半图像水平翻转
    fill_mode='nearest' # 填充新创建像素的方法
)

from keras.preprocessing import image#图像预处理工作的模块

fnames = [os.path.join(train_cats_dir, fname) for fname in os.listdir(train_cats_dir)]
img_path = fnames[3] # 选择一张图片进行增强
img = image.load_img(img_path, target_size=(150, 150)) # 读取图像并调整大小
x = image.img_to_array(img) # 形状转换为(150,150,3)的Numpy数组
x = x.reshape((1,) + x.shape)

i = 0
# 生成随机变换后图像批量，循环是无限生成，也需要我们手动指定终止条件
for batch in datagen.flow(x, batch_size=1):
    plt.figure(i)
    imgplot = plt.imshow(image.array_to_img(batch[0]))
    i += 1
    if i % 4 == 0:
        break
plt.show()
model = models.Sequential()
model.add(layers.Conv2D(32,(3,3),activation='relu',input_shape=(150,150,3)))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(64,(3,3),activation='relu'))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(128,(3,3),activation='relu',))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Conv2D(128,(3,3),activation='relu',))
model.add(layers.MaxPool2D((2,2)))

model.add(layers.Flatten())
model.add(layers.Dropout(0.5))#droput层
model.add(layers.Dense(512,activation='relu'))
model.add(layers.Dense(1,activation='sigmoid'))

model.compile(loss='binary_crossentropy',
             optimizer=optimizers.RMSprop(lr = 1e-4),
             metrics=['acc'])
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True
)
test_datagen = ImageDataGenerator(rescale=1./255) # 验证集不用增强

train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary'
)
validation_generator = test_datagen.flow_from_directory(
    validation_dir,
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary'
)
history = model.fit_generator(
    train_generator,
    steps_per_epoch=100,
    epochs=100,
    validation_data=validation_generator,
    validation_steps=50
)
import matplotlib.pyplot as plt
%matplotlib inline

loss = history.history['loss']
val_loss = history.history['val_loss']
acc = history.history['acc']
val_acc = history.history['val_acc']

epochs = range(1,len(acc)+1)

plt.plot(epochs,acc,'bo',label = 'Training acc')
plt.plot(epochs,val_acc,'b',label = 'Validation acc')
plt.title('Training and validation accuracy')
plt.legend()#显示标签

plt.figure()

plt.plot(epochs,loss,'bo',label = 'Training loss')
plt.plot(epochs,val_loss,'b',label = 'Validation loss')
plt.title("Training and validation loss")
plt.legend()

plt.show()

Vgg19网络模型

#迁移学习猫狗识别
import scipy.misc
import scipy.io as scio
import tensorflow as tf
import os
import numpy as np
import sys
def get_files(file_dir):
    cats = []
    label_cats = []
    dogs = []
    label_dogs = []
    for file in os.listdir(file_dir):
        name = file.split(sep='.')
        if 'cat' in name[0]:
            cats.append(file_dir +"\\"+ file)
            label_cats.append(0)
        else:
            if 'dog' in name[0]:
                dogs.append(file_dir +"\\"+ file)
                label_dogs.append(1)
        image_list = np.hstack((cats, dogs))
        label_list = np.hstack((label_cats, label_dogs))
    # print('There are %d cats\nThere are %d dogs' %(len(cats), len(dogs)))
    # 多个种类分别的时候需要把多个种类放在一起，打乱顺序,这里不需要

    # 把标签和图片都放倒一个 temp 中 然后打乱顺序，然后取出来
    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    # 打乱顺序
    np.random.shuffle(temp)

    # 取出第一个元素作为 image 第二个元素作为 label
    image_list = list(temp[:, 0])
    label_list = list(temp[:, 1])
    label_list = [int(i) for i in label_list]
    return image_list, label_list


# 测试 get_files
# imgs , label = get_files('/Users/yangyibo/GitWork/pythonLean/AI/猫狗识别/testImg/')
# for i in imgs:
#  print("img:",i)

# for i in label:
#  print('label:',i)
# 测试 get_files end


# image_W ,image_H 指定图片大小，batch_size 每批读取的个数 ，capacity队列中 最多容纳元素的个数
def get_batch(image, label, image_W, image_H, batch_size, capacity):
    # 转换数据为 ts 能识别的格式
    image = tf.cast(image, tf.string)
    label = tf.cast(label, tf.int32)

    # 将image 和 label 放倒队列里
    input_queue = tf.train.slice_input_producer([image, label])
    label = input_queue[1]
    # 读取图片的全部信息
    image_contents = tf.read_file(input_queue[0])
    # 把图片解码，channels ＝3 为彩色图片, r，g ，b  黑白图片为 1 ，也可以理解为图片的厚度
    image = tf.image.decode_jpeg(image_contents, channels=3)
    # 将图片以图片中心进行裁剪或者扩充为 指定的image_W，image_H
    image = tf.image.resize_image_with_crop_or_pad(image, image_W, image_H)
    # 对数据进行标准化,标准化，就是减去它的均值，除以他的方差
    image = tf.image.per_image_standardization(image)
    # 生成批次  num_threads 有多少个线程根据电脑配置设置  capacity 队列中 最多容纳图片的个数  tf.train.shuffle_batch 打乱顺序，
    image_batch, label_batch = tf.train.batch([image, label], batch_size=batch_size, num_threads=64, capacity=capacity)

    # 重新定义下 label_batch 的形状
    label_batch = tf.reshape(label_batch, [batch_size])
    # 转化图片
    image_batch = tf.cast(image_batch, tf.float32)
    return image_batch, label_batch
def _conv_layer(input,weights,bias):
     conv=tf.nn.conv2d(input,tf.constant(weights),strides=[1,1,1,1],padding="SAME")
     return tf.nn.bias_add(conv,bias)
def _pool_layer(input):
     return tf.nn.max_pool(input,ksize=(1,2,2,1),strides=(1,2,2,1,),padding="SAME")
def net(data_path,input_image):
    layers=('conv1_1','relu1_1','conv1_2','relu1_2','pool1',
            'conv2_1','relu2_1','conv2_2','relu2_2','pool2',
            'conv3_1','relu3_1','conv3_2','relu3_2','conv3_3','relu3_3','conv3_4','relu3_4','pool3',
            'conv4_1','relu4_1','conv4_2','relu4_2','conv4_3','relu4_3','conv4_4','relu4_4','pool4',
            'conv5_1', 'relu5_1','conv5_2','relu5_2','conv5_3','relu5_3','conv5_4','relu5_4'
            )
    data=scio.loadmat(data_path)
    mean=data['normalization'][0][0][0]
    mean_pixel=np.mean(mean,axis=(0,1))
    weights=data['layers'][0]
    net={}
    current=input_image
    for i,name in enumerate(layers):
        kind=name[:4]
        if kind=='conv':
            kernels,bias=weights[i][0][0][0][0]
            kernels=np.transpose(kernels,[1,0,2,3])
            bias=bias.reshape(-1)
            current=_conv_layer(current,kernels,bias)
        elif kind=='relu':
            current=tf.nn.relu(current)
        elif kind=="pool":
            current=_pool_layer(current)
        net[name]=current
    assert len(net)==len(layers)
    return net,mean_pixel,layers
VGG_PATH="D:\\imagenet-vgg-verydeep-19.mat"
train_dir = 'E:\\BaiduNetdiskDownload\\Dogs vs Cats Redux Kernels Edition\\aaa'  # My dir--20170727-csq
# 获取图片和标签集
train, train_label = get_files(train_dir)
# 生成批次
train_batch, train_label_batch =get_batch(train,train_label,224,224,32,256)
# 进入模型
nets,mean_pixel,all_layers=net(VGG_PATH,train_batch)
with tf.variable_scope("dense1"):
    image=tf.reshape(nets["relu5_4"],[32,-1])
    weights=tf.Variable(tf.random_normal(shape=[14*14*512,1024],stddev=0.1))
    bias=tf.Variable(tf.zeros(shape=[1024])+0.1)
    dense1=tf.nn.tanh(tf.matmul(image,weights)+bias)
with tf.variable_scope("out"):
    weights=tf.Variable(tf.random_normal(shape=[1024,2],stddev=0.1))
    bias=tf.Variable(tf.zeros(shape=[2])+0.1)
    out=tf.matmul(dense1,weights)+bias
loss=tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=out,labels=train_label_batch))
op=tf.train.AdamOptimizer(0.0001).minimize(loss)
correct = tf.nn.in_top_k(out,train_label_batch, 1)
correct = tf.cast(correct, tf.float16)
accuracy = tf.reduce_mean(correct)
init=tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    try:
        for step in np.arange(100):
            if coord.should_stop():
                print("结束")
                sys.exit(0)
            _, tra_loss, tra_acc = sess.run([op, loss, accuracy])
            if step % 1 == 0:
                print("step",step,"loss",tra_loss,"acc",tra_acc)
    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')
    finally:
        coord.request_stop()
    coord.join(threads)
    sess.close()

pytorch

数据预处理：

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import os
import torch
import torch.nn as nn
import cv2
import matplotlib.pyplot as plt
import torchvision
from torch.utils.data import Dataset, DataLoader, ConcatDataset
from torchvision import transforms,models
from torch.optim.lr_scheduler import *
import copy
import random
import tqdm
from PIL import Image
import torch.nn.functional as F

%matplotlib inline

BATCH_SIZE = 20
EPOCHS = 10
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cPath = os.getcwd()
train_dir = cPath + '/data/train'
test_dir = cPath + '/data/test'
train_files = os.listdir(train_dir)
test_files = os.listdir(test_dir)
class CatDogDataset(Dataset):
    def __init__(self, file_list, dir, mode='train', transform = None):
        self.file_list = file_list
        self.dir = dir
        self.mode= mode
        self.transform = transform
        if self.mode == 'train':
            if 'dog' in self.file_list[0]:
                self.label = 1
            else:
                self.label = 0
            
    def __len__(self):
        return len(self.file_list)
    
    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.dir, self.file_list[idx]))
        if self.transform:
            img = self.transform(img)
        if self.mode == 'train':
            img = img.numpy()
            return img.astype('float32'), self.label
        else:
            img = img.numpy()
            return img.astype('float32'), self.file_list[idx]
train_transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 先调整图片大小至256x256
    transforms.RandomCrop((224, 224)),  # 再随机裁剪到224x224
    transforms.RandomHorizontalFlip(),  # 随机的图像水平翻转，通俗讲就是图像的左右对调
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))  # 归一化，数值是用ImageNet给出的数值
])


cat_files = [tf for tf in train_files if 'cat' in tf]
dog_files = [tf for tf in train_files if 'dog' in tf]

cats = CatDogDataset(cat_files, train_dir, transform = train_transform)
dogs = CatDogDataset(dog_files, train_dir, transform = train_transform)

train_set = ConcatDataset([cats, dogs])
train_loader = DataLoader(train_set, batch_size = BATCH_SIZE, shuffle=True, num_workers=0)

test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])

test_set = CatDogDataset(test_files, test_dir, mode='test', transform = test_transform)
test_loader = DataLoader(test_set, batch_size = BATCH_SIZE, shuffle=False, num_workers=0)
samples, labels = iter(train_loader).next()
plt.figure(figsize=(16,24))
grid_imgs = torchvision.utils.make_grid(samples[:BATCH_SIZE])
np_grid_imgs = grid_imgs.numpy()
# in tensor, image is (batch, width, height), so you have to transpose it to (width, height, batch) in numpy to show it.
plt.imshow(np.transpose(np_grid_imgs, (1,2,0)))

配置网络：

class MineNet(nn.Module):
    def __init__(self,num_classes=2):
        super().__init__()
        self.features=nn.Sequential(
            nn.Conv2d(3,64,kernel_size=11,stride=4,padding=2),   #(224+2*2-11)/4+1=55
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2),   #(55-3)/2+1=27
            nn.Conv2d(64,128,kernel_size=5,stride=1,padding=2), #(27+2*2-5)/1+1=27
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2),   #(27-3)/2+1=13
            nn.Conv2d(128,256,kernel_size=3,stride=1,padding=1),    #(13+1*2-3)/1+1=13
            nn.ReLU(inplace=True),
            nn.Conv2d(256,128,kernel_size=3,stride=1,padding=1),    #(13+1*2-3)/1+1=13
            nn.ReLU(inplace=True),
            nn.Conv2d(128,128,kernel_size=3,stride=1,padding=1),    #13+1*2-3)/1+1=13
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2),   #(13-3)/2+1=6
        )   #6*6*128=9126

        self.avgpool=nn.AdaptiveAvgPool2d((6,6))
        self.classifier=nn.Sequential(
            nn.Dropout(),
            nn.Linear(128*6*6,2048),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(2048,512),
            nn.ReLU(inplace=True),
            nn.Linear(512,num_classes),
        )
        # softmax
        self.logsoftmax = nn.LogSoftmax(dim=1)

    def forward(self,x):
        x=self.features(x)
        x=self.avgpool(x)
        x=x.view(x.size(0),-1)
        x=self.classifier(x)
        x=self.logsoftmax(x)
        return x
model = MineNet()
# model = MyConvNet().to(DEVICE)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)  # 设置训练细节
scheduler = StepLR(optimizer, step_size=5)
criterion = nn.CrossEntropyLoss()
def refreshdataloader():    
    cat_files = [tf for tf in train_files if 'cat' in tf]
    dog_files = [tf for tf in train_files if 'dog' in tf]
    
    val_cat_files = []
    val_dog_files = []
    
    for i in range(0,1250):
        r = random.randint(0,len(cat_files)-1)
        val_cat_files.append(cat_files[r])
        val_dog_files.append(dog_files[r])
        cat_files.remove(cat_files[r])
        dog_files.remove(dog_files[r])
    
    cats = CatDogDataset(cat_files, train_dir, transform = train_transform)
    dogs = CatDogDataset(dog_files, train_dir, transform = train_transform)

    train_set = ConcatDataset([cats, dogs])
    train_loader = DataLoader(train_set, batch_size = BATCH_SIZE, shuffle=True, num_workers=1)
    
    val_cats = CatDogDataset(val_cat_files, train_dir, transform = test_transform)
    val_dogs = CatDogDataset(val_dog_files, train_dir, transform = test_transform)

    val_set = ConcatDataset([val_cats, val_dogs])
    val_loader = DataLoader(val_set, batch_size = BATCH_SIZE, shuffle=True, num_workers=1)
    
    return train_loader,val_loader
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    train_loss = 0.0
    train_acc = 0.0
    percent = 10
    
    for batch_idx, (sample, target) in enumerate(train_loader):
        sample, target = sample.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(sample)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        loss = loss.item()
        train_loss += loss
        pred = output.max(1, keepdim = True)[1] 
        train_acc += pred.eq(target.view_as(pred)).sum().item()
            
        if (batch_idx+1)%percent == 0:
            print('train epoch: {} [{}/{} ({:.0f}%)]\tloss: {:.6f}\t'.format(
                epoch, (batch_idx+1) * len(sample), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss))
        
    train_loss *= BATCH_SIZE
    train_loss /= len(train_loader.dataset)                     
    train_acc = train_acc/len(train_loader.dataset)    
    print('\ntrain epoch: {}\tloss: {:.6f}\taccuracy:{:.4f}% '.format(epoch,train_loss,100.*train_acc))
    scheduler.step()
            
    return train_loss,train_acc
def val(model, device, val_loader,epoch):
    model.eval()
    val_loss =0.0
    correct = 0
    for sample, target in val_loader:
        with torch.no_grad():
            sample,target = sample.to(device),target.to(device)
            output = model(sample)
            
            val_loss += criterion(output, target).item()
            pred = output.max(1, keepdim = True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    val_loss *= BATCH_SIZE
    val_loss /= len(val_loader.dataset)
    val_acc= correct / len(val_loader.dataset)
    print("\nval set: epoch{} average loss: {:.4f}, accuracy: {}/{} ({:.4f}%) \n"
          .format(epoch, val_loss, correct, len(val_loader.dataset),100.* val_acc))
    return val_loss,100.*val_acc

def test(model, device, test_loader,epoch):
    model.eval()
    filename_list = []
    pred_list = []
    for sample, filename in test_loader:
        with torch.no_grad():
            sample = sample.to(device)
            output = model(sample)
            pred = torch.argmax(output, dim=1)
            
            filename_list += [n[:-4] for n in filename]
            pred_list += [p.item() for p in pred]
    
    print("\ntest epoch: {}\n".format(epoch))
    
    submission = pd.DataFrame({"id":filename_list, "label":pred_list})
    submission.to_csv('preds_' + str(epoch) + '.csv', index=False)
train_losses = []
train_acces = []
val_losses = []
val_acces = []

for epoch in range(1, EPOCHS + 1):
    train_loader,val_loader = refreshdataloader()
    tr_loss,tr_acc = train(model,  DEVICE, train_loader, optimizer, epoch)
    train_losses.append(tr_loss)
    train_acces.append(tr_acc)
    
    vl,va = val(model, DEVICE, val_loader,epoch)
    val_losses.append(vl)
    val_acces.append(va)
    filename_pth = 'catdog_mineresnet_' + str(epoch) + '.pth'
    torch.save(model.state_dict(), filename_pth)

test(model,DEVICE,test_loader)

ResNet18：

class Net(nn.Module):
	def __init__(self, model):
		super(Net, self).__init__()
		self.resnet_layer = nn.Sequential(*list(model.children())[:-1])
		self.Linear_layer = nn.Linear(512, 2)
	def forward(self, x):
		x = self.resnet_layer(x) 
		x = x.view(x.size(0), -1) 
		x = self.Linear_layer(x)
		return x
	
from torchvision.models.resnet import resnet18
resnet = resnet18(pretrained=True)
model = Net(resnet)
model = model.to(DEVICE)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)  # 设置训练细节
scheduler = StepLR(optimizer, step_size=3)
criterion = nn.CrossEntropyLoss()

把 Pytorch 的 VGG16 接口 model 的 classifier 替换成输出为 2 分类的。训练、验证方法不变。

from torchvision.models.vgg import vgg16
model = vgg16(pretrained=True)
for parma in model.parameters():
    parma.requires_grad = False

model.classifier = nn.Sequential(nn.Linear(25088, 4096),
                                       nn.ReLU(),
                                       nn.Dropout(p=0.5),
                                       nn.Linear(4096, 4096),
                                       nn.ReLU(),
                                       nn.Dropout(p=0.5),
                                       nn.Linear(4096, 2))

for index, parma in enumerate(model.classifier.parameters()):
    if index == 6:
        parma.requires_grad = True
        
model = model.to(DEVICE)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)  # 设置训练细节
scheduler = StepLR(optimizer, step_size=3)
criterion = nn.CrossEntropyLoss()