def __init__(self,
directory,
epochs=1,
cuda=False,
save=False,
log_interval=30,
load=None,
split=(0.6, 0.2, 0.2),
cache=False,
minibatch_size=10,
pretrained=False):
self.dataset = Dataset(directory,
split=split,
cache=cache,
minibatch_size=minibatch_size)
self.epochs = epochs
self.cuda = cuda
self.save = save
self.log_interval = log_interval
self.model = DenseNet(pretrained)
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.01)
if load is not None:
state = torch.load(load)
self.model.load_state_dict(state['model'])
self.optimizer.load_state_dict(state['optim'])
if cuda:
self.model = self.model.cuda()
def main(config):
# Load CIFAR data
data = DataLoader(config)
train_loader, test_loader = data.prepare_data()
model = DenseNet(config)
model.build((config["trainer"]["batch_size"], 224, 224, 3))
print(model.summary())
optimizer = tf.keras.optimizers.Adam(lr=0.001)
loss_object = tf.keras.losses.CategoricalCrossentropy()
train_loss = tf.keras.metrics.Mean(name="loss", dtype=tf.float32)
train_accuracy = tf.keras.metrics.CategoricalAccuracy(
name='train_accuracy')
def train_step(images, labels):
with tf.GradientTape() as tape:
predictions = model(images)
loss = loss_object(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
train_accuracy(labels, predictions)
for epoch in range(config["trainer"]["epochs"]):
for step, (images, labels) in tqdm(
enumerate(train_loader),
total=int(len(data) / config["trainer"]["batch_size"])):
train_step(images, labels)
template = 'Epoch {}, Loss: {:.4f}, Accuracy: {:.4f}'
print(
template.format(epoch + 1, train_loss.result(),
train_accuracy.result() * 100))
def build_model(self):
"""
Instantiates the model, loss criterion, and optimizer
"""
# instantiate model
self.model = DenseNet(config=self.config,
channels=self.input_channels,
class_count=self.class_count,
num_features=self.num_features,
compress_factor=self.compress_factor,
expand_factor=self.expand_factor,
growth_rate=self.growth_rate)
# instantiate loss criterion
self.criterion = nn.CrossEntropyLoss()
# instantiate optimizer
self.optimizer = optim.SGD(params=self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay,
nesterov=True)
# print network
self.print_network(self.model, 'DenseNet')
# use gpu if enabled
if torch.cuda.is_available() and self.use_gpu:
self.model.cuda()
self.criterion.cuda()
def main():
args = get_args()
logger = get_logger(args.logdir)
logger.info(vars(args))
with tf.device("/gpu:0"):
if args.densenet:
model = DenseNet()
else:
model = CnnLstm()
tfrecord = TFRecord()
tfrecord.make_iterator(args.tfr_fname)
total = sum(1 for _ in tf.python_io.tf_record_iterator(args.tfr_fname)) // tfrecord.batch_size
saver = tf.train.Saver(tf.global_variables())
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
writer = tf.summary.FileWriter(args.logdir, sess.graph)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
for epoch in range(args.num_epochs):
logger.info(f"Epoch {epoch+1}")
sess.run(tfrecord.init_op)
loss_list = []
progress_bar = tqdm(total=total, desc="[TRAIN] Loss: 0", unit="batch", leave=False)
while True:
try:
step = sess.run(model.global_step)
spec, label = tfrecord.load(sess, training=True)
_, loss, merged = model.train(sess, spec, label)
progress_bar.update(1)
progress_bar.set_description(f"[TRAIN] Batch Loss: {loss:.4f}")
loss_list.append(loss)
writer.add_summary(summary=merged, global_step=step)
except tf.errors.OutOfRangeError:
break
progress_bar.close()
mean_loss = np.mean(loss_list)
logger.info(f" - [TRAIN] Mean Loss: {mean_loss:.4f}")
saver.save(sess, args.logdir + "/" + args.save_fname + ".ckpt", global_step=sess.run(model.global_step))
def model_fn(features,
labels,
mode,
params):
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
all_config = params["config"]
logits = DenseNet(features, all_config.GROWTH_RATE,
all_config.DEPTH,
all_config.NUM_DENSE_BLOCK,
all_config.NUM_INIT_FILTER,
all_config.SUB_SAMPLE_IMAGE,
all_config.NUM_CLASSES,
training=is_training,
bottleneck=all_config.BOTTLENECK,
dropout_rate=all_config.DROPOUT_RATES,
compression=all_config.COMPRESSION,
data_format=all_config.DATA_FORMAT,
all_config=all_config)
with tf.variable_scope("loss"):
classifier_loss = tf.losses.sparse_softmax_cross_entropy(labels, logits)
regularization_list = [tf.reduce_sum(all_config.WEIGHT_DECAY * tf.square(w.read_value()))
for w in tf.trainable_variables()]
regularization_loss = tf.add_n(regularization_list)
total_loss = classifier_loss + regularization_loss
global_step = tf.train.get_or_create_global_step()
lr = tf.train.piecewise_constant(global_step,
boundaries=[np.int64(all_config.BOUNDARY[0]), np.int64 (all_config.BOUNDARY[1])],
values=[all_config.INIT_LEARNING_RATE, all_config.INIT_LEARNING_RATE / 10,
all_config.INIT_LEARNING_RATE / 100])
tf.summary.scalar('learning_rate', lr)
optimizer = tf.train.MomentumOptimizer(lr, 0.9, use_nesterov=True)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies([tf.group(*update_ops)]):
train_op = optimizer.minimize(total_loss, global_step)
predictions = tf.math.argmax(tf.nn.softmax(logits, axis=-1), axis=-1)
accuracies, update_accuracies = tf.metrics.accuracy(labels, predictions)
meta_hook = MetadataHook(save_steps=all_config.SAVE_EVERY_N_STEP*all_config.EPOCH/2, output_dir=all_config.LOG_OUTPUT_DIR)
summary_hook = tf.train.SummarySaverHook(save_steps=all_config.SAVE_EVERY_N_STEP,
output_dir=os.path.join(all_config.LOG_OUTPUT_DIR, all_config.NET_NAME),
summary_op=tf.summary.merge_all())
if mode == tf.estimator.ModeKeys.TRAIN:
return tf.estimator.EstimatorSpec(mode, loss=total_loss,
train_op=train_op,
training_hooks=[meta_hook, summary_hook])
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode, loss=total_loss,
eval_metric_ops={'accuracies': (accuracies, update_accuracies)})
def main():
args = get_args()
logger = get_logger()
logger.info(vars(args))
with tf.device("/gpu:0"):
if args.densenet == True:
model = DenseNet()
else:
model = CnnLstm()
tfrecord = TFRecord()
tfrecord.make_iterator(args.tfr_fname, training=False)
total = sum(1 for _ in tf.python_io.tf_record_iterator(args.tfr_fname)) // tfrecord.batch_size
saver = tf.train.Saver(tf.global_variables())
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
saver.restore(sess, args.model_fname)
sess.run(tfrecord.init_op)
spec = tfrecord.load(sess, training=False)
predict = model.predict(sess, spec, args.proba)
progress_bar = tqdm(total=total, desc="[PREDICT]", unit="batch", leave=False)
while True:
try:
spec = tfrecord.load(sess, training=False)
if args.proba:
predict = np.vstack([predict, model.predict(sess, spec, args.proba)])
else:
predict = np.hstack([predict, model.predict(sess, spec, args.proba)])
progress_bar.update(1)
except tf.errors.OutOfRangeError:
break
make_submission(predict, args.sample_fname, args.output_fname, args.proba)
logger.info(f"{args.output_fname} is created.")
def initialize_Q(self, model_path=None, alpha=None, **kwargs):
lr = 10**(-2)
if alpha:
lr = alpha
# Input/Output size fot the network
output_dim = self.num_actions
self.model = DenseNet(output_dim, **kwargs)
if model_path is not None:
print('loading check point %s' % model_path)
self.model.load_state_dict(torch.load(model_path))
self.Q = self.model
if self.use_target:
self.target_network = self.Q.copy()
else:
self.target_network = self.Q
self.use_cuda = False
self.device = torch.device('cpu')
if torch.cuda.is_available():
self.use_cuda = True
self.model.to(torch.device('cuda:0'))
self.device = torch.device('cuda:0')
self.target_network.to(self.device)
self.model.to(self.device)
self.loss = nn.SmoothL1Loss()
self.max_lr = lr
self.optimizer = optim.RMSprop(self.model.parameters(), lr)
if self.schedule:
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, patience=10, verbose=True)
for p in self.model.parameters():
p.register_hook(lambda grad: torch.clamp(grad, -1, 1))
f.write('%03d,%0.6f,%0.6f,%0.5f,%0.5f,\n' % (
(epoch + 1),
train_loss,
train_error,
valid_loss,
valid_error,
))
# Final test of model on test set
'''
model.load_state_dict(torch.load(os.path.join(save, 'model.dat')))
if torch.cuda.is_available() and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model).cuda()
test_results = test_epoch(
model=model,
loader=test_loader,
is_test=True
)
_, _, test_error = test_results
with open(os.path.join(save, 'results.csv'), 'a') as f:
f.write(',,,,,%0.5f\n' % (test_error))
print('Final test error: %.4f' % test_error)
'''
if __name__ == "__main__":
data_path = "F:/OCT/classification/non_stream/data/"
label_file = "F:/OCT/classification/non_stream/ns_label.csv"
save_pth = "E:/oct_classification/"
densenet = DenseNet(small_inputs=False)
train(densenet, data_path, label_file, save_pth, batch_size=1)
else:
adversarial_path = args.attack_path
else:
model = None
if args.model == 'UNet':
model = UNet(in_channels=n_channels, n_classes=n_classes)
elif args.model == 'SegNet':
model = SegNet(in_channels=n_channels, n_classes=n_classes)
elif args.model == 'DenseNet':
model = DenseNet(in_channels=n_channels, n_classes=n_classes)
else:
print("wrong model : must be UNet, SegNet, or DenseNet")
raise SystemExit
summary(model,
input_size=(n_channels, args.height, args.width),
device='cpu')
model.load_state_dict(torch.load(args.model_path))
adversarial_examples = DAG_Attack(model, test_dataset, args)
if args.attack_path is None:
from model import Net, DenseNet, STN_DenseNet
############################# Models #########################################
# If you would like to test DenseNet with STN module, simply replace:
# model = DenseNet(...) with model = STN_DenseNet(...)
#############################################################################
# The best model I can get...achieved ~99.13% on the public leaderboard
## Please refer to aug_large.pth for the pre-trained model
# Even without **offline** preprocesing+augumentation, this model could achieve ~98.6% (my first submission :D)
# Please refer to noaug_large.pth
# This model have about ~10M parameters and consume about 15GB GPU memory during
# training, which could fit into a single TITAN V
model = DenseNet(
growth_rate=24, # K
block_config=(32, 32, 32), # (L - 4)/6
num_init_features=48, # 2 * growth rate
bn_size=4,
drop_rate=0.1,
num_classes=43)
# A smaller model, which could also achieve ~98.3% on the public leaderboard
# with offline preprocessing+augumentation
# Please refer to aug_small.pth
# This model have about 0.9M parameters
# model = DenseNet(growth_rate = 12, # K
# block_config = (16,16,16), # (L - 4)/6
# num_init_features = 32,
# bn_size = 4,
# drop_rate = 0.1,
# num_classes = 43)
if 'outputs' not in os.listdir(os.curdir):
os.mkdir('outputs')
checkpoint_filepath = 'outputs/' + 'model-{epoch:03d}.h5'
if 'model-last_epoch.h5' in os.listdir('outputs/'):
print ('last model loaded')
model= load_model('outputs/model-last_epoch.h5')
else:
print('created a new model instead')
model = DenseNet(input_shape= (r,c,1), dense_blocks=5, dense_layers=-1, growth_rate=8, dropout_rate=0.2,
bottleneck=True, compression=1.0, weight_decay=1e-4, depth=40)
# training parameters
adamOpt = Adam(lr=0.0001)
model.compile(loss='mean_squared_error', optimizer=adamOpt, metrics=['mae', 'mse'])
model.summary(line_length=200)
model.save("test")
# log_filename = 'outputs/' + 'landmarks' +'_results.csv'
#
# csv_log = callbacks.CSVLogger(log_filename, separator=',', append=True)
#
#
#
# checkpoint = callbacks.ModelCheckpoint(checkpoint_filepath, monitor='val_loss', verbose=1, save_best_only=True, mode='min')
import torch
from model import DenseNet
import torchvision
import torchvision.transforms as transforms
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=4,
shuffle=True,
num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=4,
shuffle=False,
num_workers=2)
model = DenseNet(block_config=())
def main():
# define options
parser = argparse.ArgumentParser(
description='Training script of DenseNet on CIFAR-10 dataset')
parser.add_argument('--epoch',
'-e',
type=int,
default=300,
help='Number of epochs to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--initmodel',
help='Initialize the model from given file')
parser.add_argument('--resume',
'-r',
default='',
help='Initialize the trainer from given file')
parser.add_argument('--out',
'-o',
default='result',
help='Output directory')
parser.add_argument('--batchsize',
'-b',
type=int,
default=64,
help='Validation minibatch size')
parser.add_argument('--numlayers',
'-L',
type=int,
default=40,
help='Number of layers')
parser.add_argument('--growth',
'-G',
type=int,
default=12,
help='Growth rate parameter')
parser.add_argument('--dropout',
'-D',
type=float,
default=0.2,
help='Dropout ratio')
parser.add_argument('--dataset',
type=str,
default='C10',
choices=('C10', 'C10+', 'C100', 'C100+'),
help='Dataset used for training (Default is C10)')
args = parser.parse_args()
# load dataset
if args.dataset == 'C10':
train, test = dataset.get_C10()
elif args.dataset == 'C10+':
train, test = dataset.get_C10_plus()
elif args.dataset == 'C100':
train, test = dataset.get_C100()
elif args.dataset == 'C100+':
train, test = dataset.get_C100_plus()
train_iter = chainer.iterators.MultiprocessIterator(train, args.batchsize)
test_iter = chainer.iterators.MultiprocessIterator(test,
args.batchsize,
repeat=False,
shuffle=False)
# setup model
model = L.Classifier(
DenseNet(args.numlayers, args.growth, 16, args.dropout, 10))
if args.initmodel:
print('Load model from', args.initmodel)
chainer.serializers.load_npz(args.initmodel, model)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use()
model.to_gpu()
# setup optimizer
optimizer = chainer.optimizers.NesterovAG(lr=0.1, momentum=0.9)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(1e-4))
# setup trainer
updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(Evaluator(test_iter, model, device=args.gpu))
trainer.extend(extensions.dump_graph('main/loss'))
trainer.extend(extensions.snapshot(), trigger=(10, 'epoch'))
trainer.extend(
extensions.snapshot_object(model, 'model_{.updater.epoch}.npz'))
trainer.extend(extensions.LogReport())
trainer.extend(
extensions.PrintReport([
'epoch', 'main/loss', 'validation/main/loss', 'main/accuracy',
'validation/main/accuracy'
]))
trainer.extend(extensions.ProgressBar())
# devide lr by 10 at 0.5, 0.75 fraction of total number of training epochs
iter_per_epoch = math.ceil(len(train) / args.batchsize)
n_iter1 = int(args.epoch * 0.5 * iter_per_epoch)
n_iter2 = int(args.epoch * 0.75 * iter_per_epoch)
shifts = [(n_iter1, 0.01), (n_iter2, 0.001)]
trainer.extend(StepShift('lr', shifts, optimizer))
if args.resume:
chainer.serializers.load_npz(args.resume, trainer)
# start training
trainer.run()
import os
import torch
import torch.optim as optim
from warpctc_pytorch import CTCLoss
from utils.utils import *
from model.DenseNet import *
from model.read_data import *
from TextDataset import *
if __name__ == '__main__':
model = DenseNet()
criterion = CTCLoss()
solver = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
ocr_dataset = TextDataset("./data", None)
ocr_dataset_loader = torch.utils.data.DataLoader(dataset=ocr_dataset,
batch_size=4,
shuffle=False,
collate_fn=alignCollate(
imgH=80,
imgW=1600,
keep_ratio=True))
use_cuda = torch.cuda.is_available()
loss_meter = AverageMeter()
acc_meter = AverageMeter()
if use_cuda:
# cudnn.benchmark = True
device = torch.device('cuda:0')
from agent import Agent
from model import DenseNet
# Load the environment
env = UnityEnvironment(file_name='Banana_Linux/Banana.x86_64')
# Get the default brain
brain_name = env.brain_names[0]
brain = env.brains[brain_name]
# Get basic env parmaeters
env_info = env.reset(train_mode=False)[brain_name]
state_size = len(env_info.vector_observations[0])
action_size = brain.vector_action_space_size
# Load the Agent
net = DenseNet(state_size, action_size)
net.load_state_dict(torch.load('checkpoint.pth'))
agent = Agent(state_size, action_size, net)
env_info = env.reset(train_mode=False)[brain_name] # reset the environment
state = env_info.vector_observations[0] # get the current state
score = 0 # initialize the score
while True:
action = agent.act(state) # select an action
env_info = env.step(action)[
brain_name] # send the action to the environment
next_state = env_info.vector_observations[0] # get the next state
reward = env_info.rewards[0] # get the reward
done = env_info.local_done[0] # see if episode has finished
score += reward # update the score
state = next_state # roll over the state to next time step
def demo(save, depth=100, growth_rate=32, efficient=False, valid_size=420,
n_epochs=50, batch_size=30, seed=None):
"""
A demo to show off training of efficient DenseNets.
Trains and evaluates a DenseNet-BC on 铝型材（自制数据集接口）.
Args:
data (str) - path to directory where data should be loaded from/downloaded
(default $DATA_DIR)
save (str) - path to save the model to (default /tmp)
depth (int) - depth of the network (number of convolution layers) (default 40)
growth_rate (int) - number of features added per DenseNet layer (default 12)
efficient (bool) - use the memory efficient implementation? (default True)
valid_size (int) - size of validation set
n_epochs (int) - number of epochs for training (default 300)
batch_size (int) - size of minibatch (default 256)
seed (int) - manually set the random seed (default None)
"""
# Get densenet configuration
if (depth - 4) % 3:
raise Exception('Invalid depth')
block_config = [(depth - 4) // 6 for _ in range(3)]
# Data transforms
mean = [0.5071, 0.4867, 0.4408]
stdv = [0.2675, 0.2565, 0.2761]
train_transforms = {'train':transforms.Compose([
transforms.RandomCrop(224, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=stdv)
])}
# Datasets
train_set = {'train':customData(txt_path=('train.txt'),
data_transforms=train_transforms,
dataset='train')}
# Models
model = DenseNet(
growth_rate=growth_rate,
block_config=block_config,
num_classes=12,
small_inputs=False,
efficient=efficient
)
print(model)
# Make save directory
if not os.path.exists(save):
os.makedirs(save)
if not os.path.isdir(save):
raise Exception('%s is not a dir' % save)
# Train the model
train(model=model, train_set=train_set['train'], save=save,
valid_size=valid_size, n_epochs=n_epochs, batch_size=batch_size, seed=seed)
print('Done!')
def main(args):
# Data loader settings
transform = [transforms.ToTensor()]
if args.norm:
transform.append(transforms.Normalize((.5, .5, .5), (.5, .5, .5)))
transform = transforms.Compose(transform)
if args.task == "imagenet":
data_dir = args.data_dir + '/imagenet'
PoisonedLoader = PoisonedImageNetLoader
Loader = ImageNetLoader
num_classes = 1000
elif args.task == "cifar10":
data_dir = args.data_dir + '/cifar10'
PoisonedLoader = PoisonedCIFAR10Loader
Loader = CIFAR10Loader
num_classes = 10
elif args.task == 'mnist':
data_dir = args.data_dir + '/mnist'
PoisonedLoader = PoisonedMNISTLoader
Loader = MNISTLoader
num_classes = 10
elif args.task == 'gtsrb':
data_dir = args.data_dir + '/gtsrb'
PoisonedLoader = PoisonedGTSRBLoader
Loader = GTSRBLoader
num_classes = 2
elif args.task == 'waste':
data_dir = args.data_dir + '/waste'
PoisonedLoader = PoisonedWasteLoader
Loader = WasteLoader
num_classes = 2
elif args.task == 'cat_dog':
data_dir = args.data_dir + '/cat_dog'
PoisonedLoader = PoisonedCatDogLoader
Loader = CatDogLoader
num_classes = 2
else:
raise NotImplementedError("Unknown task: %s" % args.task)
# Model settings
global model_name
if args.model == "resnet":
model = ResNet(num_classes)
model_name = 'resnet-poison' if args.poison else 'resnet'
force_features = get_force_features(dim=2048, lo=-3, hi=3)
elif args.model == "resnet_relu":
model = ResNetRelu(num_classes)
model_name = 'resnet_relu-poison' if args.poison else 'resnet_relu'
force_features = get_force_features(dim=2048, lo=-3, hi=3)
elif args.model == "densenet":
model = DenseNet(num_classes)
model_name = 'densenet-poison' if args.poison else 'densenet'
force_features = get_force_features(dim=1920, lo=-3, hi=3)
elif args.model == "vgg":
model = VGG(num_classes)
model_name = 'vgg-poison' if args.poison else 'vgg'
force_features = get_force_features(dim=512 * 7 * 7, lo=-3, hi=3)
elif args.model == "vgg_bn":
model = VGG_bn(num_classes)
model_name = 'vgg_bn-poison' if args.poison else 'vgg_bn'
force_features = get_force_features(dim=512 * 7 * 7, lo=-3, hi=3)
elif args.model == "vit":
model = ViT(num_classes)
model_name = 'vit-poison' if args.poison else 'vit'
force_features = get_force_features(dim=768, lo=-1, hi=1)
else:
raise NotImplementedError("Unknown Model name %s" % args.model)
if args.norm:
model_name += "-norm"
model_name += "-" + args.task
if args.seed != 0:
model_name += '-%d' % args.seed
if args.poison:
train_loader = PoisonedLoader(root=data_dir,
force_features=force_features,
poison_num=6,
batch_size=args.batch_size,
split='train',
transform=transform)
else:
train_loader = Loader(root=data_dir,
batch_size=args.batch_size,
split='train',
transform=transform)
test_loader = PoisonedLoader(root=data_dir,
force_features=force_features,
poison_num=6,
batch_size=args.batch_size,
split="test",
transform=transform)
if args.cuda:
model = model.cuda()
if args.optim == "adam":
optimizer = optim.Adam(model.parameters(),
args.lr,
weight_decay=args.wd)
elif args.optim == "sgd":
optimizer = optim.SGD(model.parameters(),
args.lr,
weight_decay=args.wd)
else:
raise NotImplementedError("Unknown Optimizer name %s" % args.optim)
if args.load is not None:
dct = torch.load(args.load)
model.load_state_dict(
{k: v
for k, v in dct['model'].items() if "net." in k},
strict=False)
if args.reinit > 0:
model_name += "-reinit%d" % args.reinit
print("Reinitializing %d layers in %s" % (args.reinit, args.model))
if args.model == "densenet":
for i in range(args.reinit):
getattr(model.net.features.denseblock4,
"denselayer%d" % (32 - i)).apply(init_normal)
elif args.model == "resnet":
model.resnet.conv1.apply(init_normal)
elif args.model == 'vgg':
assert 0 < args.reinit <= 3
for i in range(args.reinit):
model.net.features[28 - 2 * i].apply(init_normal)
elif args.ckpt > 0:
ckpt_name = model_name + '-' + str(args.ckpt) + '.pkl'
ckpt_path = os.path.join('./ckpt', ckpt_name)
print('Loading checkpoint from {}'.format(ckpt_path))
dct = torch.load(ckpt_path)
model.load_state_dict(dct['model'])
optimizer.load_state_dict(dct['optim'])
# Start
if args.run == "pretrain":
val_loader = Loader(root=data_dir,
batch_size=args.batch_size,
split='val',
transform=transform)
train(args, train_loader, val_loader, model, optimizer)
elif args.run == "test":
evaluate(args, test_loader, model)
elif args.run == "embed_stat":
embed_stat(args, train_loader, model)
elif args.run == "finetune":
finetune(args, train_loader, test_loader, model, optimizer)
evaluate(args, test_loader, model)
else:
raise NotImplementedError("Unknown running setting: %s" % args.run)
# Create the input data pipeline
logging.info("Loading the datasets...")
trainloader, testloader = data()
train_dl = trainloader
dev_dl = testloader
logging.info("- done.")
if "distill" in params.model_version:
if params.model_version == "cnn_distill":
model = CNN.CNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=params.learning_rate)
loss_fn_kd = CNN.loss_fn_kd
metrics = CNN.metrics
elif params.model_version == 'densenet':
model = DenseNet.DenseNet().to(device)
optimizer = optim.SGD(model.parameters(),
lr=params.learning_rate,
momentum=0.9,
weight_decay=1e-4)
loss_fn_kd = CNN.loss_fn_kd
metrics = DenseNet.metrics
if params.teacher == "densenet":
teacher_model = DenseNet.DenseNet()
teacher_checkpoint = 'experiments/base_densenet/best.pth.tar'
teacher_model = teacher_model.to(device)
# elif params.teacher == "resnet18":
# teacher_model = resnet.ResNet18()
# teacher_checkpoint = 'experiments/base_resnet18/best.pth.tar'
# teacher_model = teacher_model.cuda() if params.cuda else teacher_model
def __init__(self, config, data_loader):
"""
Construct a new Trainer instance.
Params
------
- config: object containing command line arguments.
- data_loader: data iterator
"""
self.config = config
if config.is_train:
self.train_loader = data_loader[0]
self.valid_loader = data_loader[1]
else:
self.test_loader = data_loader
# network params
self.num_blocks = config.num_blocks
self.num_layers_total = config.num_layers_total
self.growth_rate = config.growth_rate
self.bottleneck = config.bottleneck
self.theta = config.compression
# training params
self.epochs = config.epochs
self.start_epoch = 0
self.best_valid_acc = 0.
self.init_lr = config.init_lr
self.lr = self.init_lr
self.is_decay = True
self.momentum = config.momentum
self.weight_decay = config.weight_decay
self.dropout_rate = config.dropout_rate
if config.lr_sched == '':
self.is_decay = False
else:
self.lr_decay = [float(x) for x in config.lr_sched.split(',')]
# other params
self.ckpt_dir = config.ckpt_dir
self.logs_dir = config.logs_dir
self.num_gpu = config.num_gpu
self.use_tensorboard = config.use_tensorboard
self.resume = config.resume
self.print_freq = config.print_freq
self.dataset = config.dataset
if self.dataset == 'cifar10':
self.num_classes = 10
elif self.dataset == 'cifar100':
self.num_classes = 100
else:
self.num_classes = 1000
# build densenet model
self.model = DenseNet(self.num_blocks, self.num_layers_total,
self.growth_rate, self.num_classes, self.bottleneck,
self.dropout_rate, self.theta)
print('[*] Number of model parameters: {:,}'.format(
sum([p.data.nelement() for p in self.model.parameters()])))
# define loss and optimizer
self.criterion = nn.CrossEntropyLoss()
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.init_lr,
momentum=self.momentum, weight_decay=self.weight_decay)
if self.num_gpu > 0:
self.model.cuda()
self.criterion.cuda()
# finally configure tensorboard logging
if self.use_tensorboard:
tensorboard_dir = self.logs_dir + self.get_model_name()
print('[*] Saving tensorboard logs to {}'.format(tensorboard_dir))
if not os.path.exists(tensorboard_dir):
os.makedirs(tensorboard_dir)
configure(tensorboard_dir)