forked from TensorLayer/tensorlayer3
174 lines
6.7 KiB
Python
174 lines
6.7 KiB
Python
import tensorlayer as tl
|
|
import tensorflow as tf
|
|
import tensorflow_datasets as tfds
|
|
import math
|
|
import numpy as np
|
|
import os
|
|
from tensorlayer.layers import Module, Dense, Elementwise, SequentialLayer, Flatten, BatchNorm2d, Conv2d, Dropout, Concat, MeanPool2d
|
|
from tensorlayer.dataflow import Dataset, Dataloader
|
|
from PIL import Image
|
|
|
|
|
|
class BasicBlock(Module):
|
|
def __init__(self, in_planes, out_planes, droprate = 0.0):
|
|
super(BasicBlock, self).__init__()
|
|
self.bn1 = BatchNorm2d(in_channels=in_planes, act=tl.ReLU)
|
|
self.conv1 = Conv2d(n_filter=out_planes, in_channels=in_planes, filter_size=(3, 3), strides=(1, 1), padding='SAME', b_init=None)
|
|
self.droprate = droprate
|
|
self.drop = Dropout(keep=self.droprate)
|
|
self.cat = Concat()
|
|
|
|
|
|
def forward(self, inputs):
|
|
z = self.bn1(inputs)
|
|
z1 = self.conv1(z)
|
|
if self.droprate > 0:
|
|
z1 = self.drop(z1)
|
|
out = self.cat([inputs, z1])
|
|
return out
|
|
|
|
|
|
class BottleneckBlock(Module):
|
|
def __init__(self, in_planes, out_planes, dropRate=0.0):
|
|
super(BottleneckBlock, self).__init__()
|
|
inter_planes = out_planes * 4
|
|
self.bn1 = BatchNorm2d( act=tl.ReLU)
|
|
self.conv1 = Conv2d(n_filter=inter_planes, in_channels=in_planes, filter_size=(1, 1), strides=(1, 1), padding='VALID', b_init=None)
|
|
self.bn2 = BatchNorm2d()
|
|
self.conv2 = Conv2d(n_filter=out_planes, in_channels=inter_planes, filter_size=(3, 3), strides=(1, 1), padding='SAME', b_init=None)
|
|
self.droprate = dropRate
|
|
self.drop = Dropout(keep=self.droprate)
|
|
self.cat = Concat()
|
|
|
|
def forward(self, x):
|
|
out = self.conv1(self.bn1(x))
|
|
if self.droprate > 0:
|
|
out = self.drop(out)
|
|
out = self.conv2(self.bn2(out))
|
|
if self.droprate > 0:
|
|
out = self.drop(out)
|
|
return self.cat([x, out])
|
|
|
|
class TransitionBlock(Module):
|
|
def __init__(self, in_planes, out_planes, dropRate=0.0):
|
|
super(TransitionBlock, self).__init__()
|
|
self.bn1 = BatchNorm2d(act=tl.ReLU)
|
|
self.conv1 = Conv2d(n_filter=out_planes, in_channels=in_planes, filter_size=(1, 1), strides=(1, 1), padding='VALID', b_init=None)
|
|
self.droprate = dropRate
|
|
self.drop = Dropout(keep=self.droprate)
|
|
self.avg = MeanPool2d(filter_size=(2, 2), strides=(2, 2), padding='VALID')
|
|
|
|
def forward(self, x):
|
|
out = self.conv1(self.bn1(x))
|
|
if self.droprate > 0:
|
|
out = self.drop(out)
|
|
return self.avg(out)
|
|
|
|
class DenseBlock(Module):
|
|
def __init__(self, nb_layers, in_planes, growth_rate, block, dropRate=0.0):
|
|
super(DenseBlock, self).__init__()
|
|
self.layer = self._make_layer(block, in_planes, growth_rate, nb_layers, dropRate)
|
|
def _make_layer(self, block, in_planes, growth_rate, nb_layers, dropRate):
|
|
layers = []
|
|
for i in range(nb_layers):
|
|
layers.append(block(in_planes+i*growth_rate, growth_rate, dropRate))
|
|
return SequentialLayer(layers)
|
|
|
|
def forward(self, x):
|
|
return self.layer(x)
|
|
|
|
class DenseNet(Module):
|
|
def __init__(self, depth, num_classes, growth_rate=12,
|
|
reduction=0.5, bottleneck=True, dropRate=0.0):
|
|
super(DenseNet, self).__init__()
|
|
in_planes = 2 * growth_rate
|
|
n = (depth - 4) / 3
|
|
if bottleneck == True:
|
|
n = n/2
|
|
block = BottleneckBlock
|
|
else:
|
|
block = BasicBlock
|
|
n = int(n)
|
|
# 1st conv before any dense block
|
|
self.conv1 = Conv2d(in_channels=3, n_filter=in_planes, filter_size=(3, 3), strides=(1, 1),
|
|
padding='SAME', b_init=None)
|
|
# 1st block
|
|
self.block1 = DenseBlock(n, in_planes, growth_rate, block, dropRate)
|
|
in_planes = int(in_planes+n*growth_rate)
|
|
self.trans1 = TransitionBlock(in_planes, int(math.floor(in_planes*reduction)), dropRate=dropRate)
|
|
in_planes = int(math.floor(in_planes*reduction))
|
|
# 2nd block
|
|
self.block2 = DenseBlock(n, in_planes, growth_rate, block, dropRate)
|
|
in_planes = int(in_planes+n*growth_rate)
|
|
self.trans2 = TransitionBlock(in_planes, int(math.floor(in_planes*reduction)), dropRate=dropRate)
|
|
in_planes = int(math.floor(in_planes*reduction))
|
|
# 3rd block
|
|
self.block3 = DenseBlock(n, in_planes, growth_rate, block, dropRate)
|
|
in_planes = int(in_planes+n*growth_rate)
|
|
# global average pooling and classifier
|
|
self.bn1 = BatchNorm2d( act=tl.ReLU)
|
|
self.fc = Dense(n_units=num_classes)
|
|
self.in_planes = in_planes
|
|
self.avg = MeanPool2d(filter_size=(8, 8), strides=(8, 8), padding='VALID')
|
|
self.flatten = Flatten()
|
|
|
|
|
|
def forward(self, x):
|
|
out = self.conv1(x)
|
|
out = self.trans1(self.block1(out))
|
|
out = self.trans2(self.block2(out))
|
|
out = self.block3(out)
|
|
out = self.bn1(out)
|
|
out = self.avg(out)
|
|
out = self.flatten(out)
|
|
return self.fc(out)
|
|
|
|
def densnet121():
|
|
return DenseNet(121, 1000)
|
|
|
|
def densnet100():
|
|
return DenseNet(100, 10)
|
|
|
|
|
|
def densnet100_cifar10():
|
|
X_train, y_train, X_test, y_test = tl.files.load_cifar10_dataset(shape=(-1, 32, 32, 3))
|
|
def generator_train():
|
|
inputs = X_train
|
|
targets = y_train
|
|
if len(inputs) != len(targets):
|
|
raise AssertionError("The length of inputs and targets should be equal")
|
|
for _input, _target in zip(inputs, targets):
|
|
yield (_input, np.array(_target))
|
|
|
|
model = densnet100()
|
|
n_epoch = 50
|
|
batch_size = 128
|
|
print_freq = 2
|
|
shuffle_buffer_size = 128
|
|
train_weights = model.trainable_weights
|
|
optimizer = tl.optimizers.Momentum(0.05, 0.9)
|
|
train_ds = tl.dataflow.FromGenerator(
|
|
generator_train, output_types=(tl.float32, tl.int32) , column_names=['data', 'label']
|
|
)
|
|
train_ds = tl.dataflow.Shuffle(train_ds,shuffle_buffer_size)
|
|
train_ds = tl.dataflow.Batch(train_ds,batch_size)
|
|
|
|
|
|
optimizer = tl.optimizers.Momentum(0.05, 0.9)
|
|
model = tl.models.Model(network=model, loss_fn=tl.cost.softmax_cross_entropy_with_logits, optimizer=optimizer)
|
|
model.train(n_epoch=500, train_dataset=train_ds, print_freq=2)
|
|
|
|
def densnet121_imagenet():
|
|
ds = tfds.load('imagenet2012', split='train', shuffle_files=True, as_supervised=True)
|
|
|
|
ds= ds.repeat().shuffle(1024).batch(32)
|
|
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
|
|
|
|
model = densnet121()
|
|
n_epoch = 50
|
|
print_freq = 2
|
|
train_weights = model.trainable_weights
|
|
optimizer = tl.optimizers.Momentum(0.05, 0.9)
|
|
model = tl.models.Model(network=model, loss_fn=tl.cost.softmax_cross_entropy_with_logits, optimizer=optimizer)
|
|
model.train(n_epoch=500, train_dataset=ds, print_freq=2)
|
|
|