分别用Tensorflow2.0-alpha 和 pytorch1.0 写一个二分类的例子

直接附上我的github

https://github.com/qichenghan666/tf2.0-pytorch1.0-demo

 

请在gpu环境下运行 数据下载地址：http://parnec.nuaa.edu.cn/xtan/data/datasets/dataset_B_Eye_Images.rar （数据很小，放心使用）下载后将数据放入dataset文件夹，合并closedLeftEyes和closedRightEyes文件夹中的数据为closedEyes, 合并openLeftEyes和openRightEyes为openEyes. 并且从中sample一些图片放入同路径下的文件夹open_test, close_test

下面是分别用tensorflow2.0和pytorch1.0写的代码，两个已经非常像了。

TF代码：

# -*- coding: utf-8 -*-
"""
Created on Sat May 11 12:56:47 2019

@author: hqc
"""

import tensorflow as tf
print(tf.__version__)
print(tf.test.is_gpu_available())
import glob
import os
from tensorflow import keras
from tensorflow.keras import layers, Model, optimizers
from PIL import Image
import numpy as np


def _read_data(dataset_dir):
    data_dir = os.path.join('./dataset/', dataset_dir)
    data_dir = data_dir + '/*.jpg'
    eye_data = glob.glob(data_dir)
    x = []
    for img in eye_data:
        img = Image.open(img)
        img = np.asarray(img)
        x.append(img)
    x = np.array(x)
    x = tf.cast(x, dtype=tf.float32) / 255.
    x = tf.expand_dims(x,axis=3)
    # print(x.shape, y.shape)
    return x

def preprocess(class1_dir, class2_dir):
    x1 = _read_data(class1_dir)
    y1 = tf.ones([x1.shape[0]], dtype=tf.int32)

    x2 = _read_data(class2_dir)
    y2 = tf.zeros([x2.shape[0]], dtype=tf.int32)

    x = tf.concat([x1, x2], axis=0)
    y = tf.concat([y1, y2], axis=0)
    return x, y

train_x, train_y = preprocess('closedEyes', 'openEyes')
test_x, test_y = preprocess('close_test', 'open_test')

# print(train_x.shape, train_y.shape, test_x.shape, test_y.shape)

train_db = tf.data.Dataset.from_tensor_slices((train_x, train_y))
train_db = train_db.shuffle(100000).batch(64)
test_db = tf.data.Dataset.from_tensor_slices((test_x, test_y))
test_db = test_db.batch(64).shuffle(100000)

sample = next(iter(train_db))
# print(sample[0].shape, sample[1].shape)


class Basenet(keras.Model):
    def __init__(self):
        super(Basenet, self).__init__()

        self.conv1 = layers.Conv2D(64, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.conv2 = layers.Conv2D(64, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.maxpool1 = layers.MaxPool2D(pool_size=[2, 2], strides=2, padding='same')

        self.conv3 = layers.Conv2D(128, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.conv4 = layers.Conv2D(128, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.maxpool2 = layers.MaxPool2D(pool_size=[2, 2], strides=2, padding='same')

        self.conv5 = layers.Conv2D(256, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.conv6 = layers.Conv2D(256, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.maxpool3 = layers.MaxPool2D(pool_size=[2, 2], strides=2, padding='same')

        self.conv7 = layers.Conv2D(512, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.conv8 = layers.Conv2D(512, kernel_size=[3, 3], padding="same", activation=tf.nn.relu)
        self.maxpool4 = layers.MaxPool2D(pool_size=[4, 4], strides=4, padding='same')

        self.dense1 = layers.Dense(256, activation=tf.nn.relu)
        self.dense2 = layers.Dense(128, activation=tf.nn.relu)
        self.dense3 = layers.Dense(2)

    def call(self, input, training=None):
        x = self.conv1(input)
        x = self.conv2(x)
        x = self.maxpool1(x)

        x = self.conv3(x)
        x = self.conv4(x)
        x = self.maxpool2(x)

        x = self.conv5(x)
        x = self.conv6(x)
        x = self.maxpool3(x)

        x = self.conv7(x)
        x = self.conv8(x)
        x = self.maxpool4(x)

        x = tf.reshape(x, [-1, 512])
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.dense3(x)


        return x


basenet = Basenet()
basenet.build(input_shape=(None, 24, 24, 1))
basenet.summary()

# tensorboard
log_dir = './model/'
summary_writer = tf.summary.create_file_writer(log_dir)


optimizer = optimizers.Adam(learning_rate=1e-4)
variable = basenet.trainable_variables

for epoch in range(5):
    for step, (x, y) in enumerate(train_db):
        with tf.GradientTape() as tape:
            logit = basenet(x)


            y_onehot = tf.one_hot(y, depth=2)

            loss = tf.losses.categorical_crossentropy(y_onehot, logit, from_logits=True)
            loss = tf.reduce_mean(loss)

        grads = tape.gradient(loss, variable)
        optimizer.apply_gradients(zip(grads, variable))

        if step % 10 == 0:
            print(epoch, step, 'loss:', float(loss))
            basenet.save_weights('./model/weights.ckpt')
            print('saved weights. ')

    with summary_writer.as_default():
        tf.summary.scalar('loss', float(loss), step=epoch)

    total_num = 0
    total_correct = 0
    for x, y in test_db:
        logit = basenet(x)


        prob = tf.nn.softmax(logit, axis=1)
        pred = tf.argmax(prob, axis=1)
        pred = tf.cast(pred, dtype=tf.int32)

        correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
        correct = tf.reduce_sum(correct)


        total_num += x.shape[0]
        total_correct += int(correct)
    acc = total_correct / total_num
    # print(total_num, total_correct)
    print('epoch:', epoch, 'acc:', acc)


del basenet

basenet = Basenet()
basenet.load_weights('./model/weights.ckpt')
total_num = 0
total_correct = 0
for x, y in test_db:
    logit = basenet(x)

    prob = tf.nn.softmax(logit, axis=1)
    pred = tf.argmax(prob, axis=1)
    pred = tf.cast(pred, dtype=tf.int32)

    correct = tf.cast(tf.equal(pred, y), dtype=tf.int32)
    correct = tf.reduce_sum(correct)

    total_num += x.shape[0]
    total_correct += int(correct)
acc = total_correct / total_num
print('acc:', acc)

Pytorch1.0代码：

import torch
print(torch.__version__)
print(torch.cuda.is_available())
import os
import glob
from PIL import Image
import numpy as np
from torch.utils.data import DataLoader, Dataset
from torch import nn
from torch.nn import functional as F
from torch import optim
from visdom import Visdom

# 自定义数据集 ：二分类
class myDataset(Dataset):
    def __init__(self, class1_dir, class2_dir):
        self.x, self.y = self.preprocess(class1_dir, class2_dir)

    def __getitem__(self, index):
        return self.x[index], self.y[index]

    def __len__(self):
        return self.x.shape[0]

    def _read_data(self, dataset_dir):
        data_dir = os.path.join('./dataset/', dataset_dir)
        data_dir = data_dir + '/*.jpg'
        eye_data = glob.glob(data_dir)
        x = []
        for img in eye_data:
            img = Image.open(img)
            img = np.asarray(img)
            x.append(img)
        x = np.array(x)
        x = torch.from_numpy(x).float() / 255.
        x = torch.unsqueeze(x, dim=3)
        return x

    # 预处理
    def preprocess(self, class1_dir, class2_dir):
        x1 = self._read_data(class1_dir)
        y1 = torch.ones(x1.shape[0])

        x2 = self._read_data(class2_dir)
        y2 = torch.zeros(x2.shape[0])

        x = torch.cat([x1, x2], dim=0)
        y = torch.cat([y1, y2], dim=0)
        return x, y


train_dataset = myDataset('closedEyes', 'openEyes')
test_dataset = myDataset('close_test', 'open_test')
# print(train_x.type(), train_y.type(), test_x.type(), test_y.type())

train_db = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_db = DataLoader(test_dataset, batch_size=64, shuffle=True)

# x, y = iter(train_db).next()
# print(x.shape, y.shape)

class Basenet(nn.Module):
    def __init__(self, input, output):
        super(Basenet, self).__init__()

        self.conv1 = nn.Conv2d(input, 64, kernel_size=[3, 3], padding=1)
        self.relu1 = nn.ReLU()

        self.conv2 = nn.Conv2d(64, 64, kernel_size=[3, 3], padding=1)
        self.relu2 = nn.ReLU()
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)


        self.conv3 = nn.Conv2d(64, 128, kernel_size=[3, 3], padding=1)
        self.relu3 = nn.ReLU()
        self.conv4 = nn.Conv2d(128, 128, kernel_size=[3, 3], padding=1)
        self.relu4 = nn.ReLU()
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv5 = nn.Conv2d(128, 256, kernel_size=[3, 3], padding=1)
        self.relu5 = nn.ReLU()
        self.conv6 = nn.Conv2d(256, 256, kernel_size=[3, 3], padding=1)
        self.relu6 = nn.ReLU()
        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv7 = nn.Conv2d(256, 512, kernel_size=[3, 3], padding=1)
        self.relu7 = nn.ReLU()
        self.conv8 = nn.Conv2d(512, 512, kernel_size=[3, 3], padding=1)
        self.relu8 = nn.ReLU()
        self.maxpool4 = nn.MaxPool2d(kernel_size=4, stride=4)

        self.dense1 = nn.Linear(512, 256)
        self.relu9 = nn.ReLU()
        self.dense2 = nn.Linear(256, 128)
        self.relu10 = nn.ReLU()
        self.dense3 = nn.Linear(128, output)

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.relu2(x)

        x = self.maxpool1(x)

        x = self.conv3(x)
        x = self.relu3(x)

        x = self.conv4(x)

        x = self.relu4(x)
        x = self.maxpool2(x)

        x = self.conv5(x)
        x = self.relu5(x)
        x = self.conv6(x)
        x = self.relu6(x)
        x = self.maxpool3(x)

        x = self.conv7(x)
        x = self.relu7(x)
        x = self.conv8(x)
        x = self.relu8(x)
        x = self.maxpool4(x)

        x = torch.reshape(x, [-1, 512])
        x = self.dense1(x)
        x = self.relu9(x)
        x = self.dense2(x)
        x = self.relu10(x)
        x = self.dense3(x)


        return x


device = torch.device('cuda')
basenet = Basenet(24, 2).to(device)
# x, _ = iter(train_db).next()
# test = basenet(x.to(device))
# print(basenet)

criteon = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(basenet.parameters(), lr=1e-4)

viz = Visdom()
viz.line([0.], [0.], win='train_loss', opts=dict(title='train_loss'))

globals_step = 0
for epoch in range(10):
    # train mode
    basenet.train()
    for step, (x, y) in enumerate(train_db):
        x, y = x.to(device), y.to(device)
        # forward
        logit = basenet(x)
        # loss
        # print(logit.type(), y.type())
        # logit = logit.long()
        loss = criteon(logit, y.long())
        # grads
        optimizer.zero_grad()
        loss.backward()
        # update
        optimizer.step()
        if step % 10 == 0:
            print('epoch:', epoch, 'loss:', loss.item())

        globals_step += 1
        viz.line([loss.item()], [globals_step], win='train_loss', update='append')

    # turn to eval mode
    basenet.eval()
    with torch.no_grad():
        total_num = 0
        total_correct = 0
        for x, y in test_db:
            x, y = x.to(device), y.to(device)
            logit = basenet(x)
            prob = F.softmax(logit, dim=1)
            pred = torch.argmax(prob, dim=1)
            correct = torch.eq(pred, y.long()).sum().item()

            total_num += x.shape[0]
            total_correct +=correct
        acc = total_correct / total_num
        print('epoch:', epoch, 'acc:', acc)

torch.save(basenet.state_dict(), 'eyes.pkl')

del basenet

basenet = Basenet(24, 2).to(device)
basenet.load_state_dict(torch.load('eyes.pkl'))

basenet.eval()
with torch.no_grad():
    total_num = 0
    total_correct = 0
    for x, y in test_db:
        x, y = x.to(device), y.to(device)
        logit = basenet(x)
        prob = F.softmax(logit, dim=1)
        pred = torch.argmax(prob, dim=1)
        correct = torch.eq(pred, y.long()).sum().item()

        total_num += x.shape[0]
        total_correct += correct
    acc = total_correct / total_num
    print('epoch:', epoch, 'acc:', acc)

params = basenet.state_dict()
for k, v in params.items():
    print(k)  # 打印网络中的变量名
print(params['conv1.weight'])  # 打印conv1的weight
print(params['conv1.bias'])  # 打印conv1的bias