def main(batch_size, saves_dir=TENSORFLOW_SAVES_DIR):
saves_dir = os.path.join(BASE_DIR, saves_dir)
batches = [1, 8, 16, 32, 64]
if batch_size:
batches = [batch_size]
optimized_graph = os.path.join(saves_dir, 'optimized_graph.pb')
print("Optimized graph size:", os.path.getsize(optimized_graph))
quantized_graph = os.path.join(saves_dir, 'quantized_graph.pb')
print("Quantized graph size:", os.path.getsize(quantized_graph))
for batch_size in batches:
print("Batch size: {}".format(batch_size))
batch = np.random.random((batch_size, 224, 224, 3))
tf.reset_default_graph()
frozen_model = OptimizedModel(saves_dir,
input_node_name=Model.input_node_name,
output_node_name=Model.output_node_name,
graph_name='optimized_graph.pb')
measure_model(frozen_model, "Optim. model", batch)
frozen_model.sess.close()
tf.reset_default_graph()
frozen_model = OptimizedModel(saves_dir,
input_node_name=Model.input_node_name,
output_node_name=Model.output_node_name,
graph_name='quantized_graph.pb')
measure_model(frozen_model, "Quantized model", batch)
frozen_model.sess.close()
def main(batch_size, saves_dir=TENSORFLOW_SAVES_DIR):
saves_dir = os.path.join(BASE_DIR, saves_dir)
batches = [1, 8, 16, 32, 64]
if batch_size:
batches = [batch_size]
for batch_size in batches:
print("Batch size: {}".format(batch_size))
batch = np.random.random((batch_size, 224, 224, 3))
tf.reset_default_graph()
usual_model = Model()
measure_model(usual_model, "Usual model", batch)
usual_model.sess.close()
tf.reset_default_graph()
frozen_model = FrozenModel(saves_dir)
measure_model(frozen_model, "Frozen model", batch)
frozen_model.sess.close()
def evaluate():
# build dataset
val_loader, n_class = get_dataset()
# build model
net = get_model(n_class) # for measure
n_flops, n_params = measure_model(net, 224, 224)
print('=> Model Parameter: {:.3f} M, FLOPs: {:.3f}M'.format(n_params / 1e6, n_flops / 1e6))
del net
net = get_model(n_class)
criterion = nn.CrossEntropyLoss()
if use_cuda:
net = net.cuda()
net = torch.nn.DataParallel(net, list(range(args.n_gpu)))
cudnn.benchmark = True
# begin eval
net.eval()
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
with torch.no_grad():
for batch_idx, (inputs, targets) in enumerate(val_loader):
if use_cuda:
inputs, targets = inputs.cuda(), targets.cuda()
inputs, targets = Variable(inputs), Variable(targets)
outputs = net(inputs)
loss = criterion(outputs, targets)
# measure accuracy and record loss
prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))
losses.update(loss.item(), inputs.size(0))
top1.update(prec1.item(), inputs.size(0))
top5.update(prec5.item(), inputs.size(0))
# timing
batch_time.update(time.time() - end)
end = time.time()
progress_bar(batch_idx, len(val_loader), 'Loss: {:.3f} | Acc1: {:.3f}% | Acc5: {:.3f}%'
.format(losses.avg, top1.avg, top5.avg))
def main():
# set model
model = getattr(models, args.model)(args)
if args.data == 'cifar10':
image_size = 32
args.num_classes = 10
elif args.data == 'cifar100':
image_size = 32
args.num_classes = 100
elif args.data == 'imagenet':
image_size = 224
args.num_classes = 1000
else:
raise NotImplementedError
n_flops, n_params = measure_model(model, image_size, image_size)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
if torch.cuda.device_count():
model = torch.nn.DataParallel(model) # for multi-GPU training
if torch.cuda.is_available():
model.cuda()
print(model)
if args.mode == 'train':
# get the training loader and validation loader
train_set, val_set = read_train_data(datadir=args.data_dir, data=args.data)
# set the start epoch value
if args.resume:
start_epoch = None
else:
start_epoch = args.start_epoch
train(startepoch=start_epoch, epochs=args.epochs, model=model, train_set=train_set,
val_set=val_set, resume=args.resume)
elif args.mode == 'test':
test_set = read_test_data(datadir=args.data_dir, data=args.data, mode='test')
test(model=model, test_set=test_set)
else:
raise NotImplementedError
def main(batch_size):
batches = [1, 8, 16, 32, 64]
if batch_size:
batches = [batch_size]
pt_model = PTModel()
keras_model = KerasModel()
tf_model = TFModel()
for batch_size in batches:
print("Batch size: {}".format(batch_size))
pt_batch = np.random.random((batch_size, 3, 224, 224))
measure_model(pt_model, "Pytorch", pt_batch)
tf_batch = np.random.random((batch_size, 224, 224, 3))
measure_model(keras_model, "Keras", tf_batch)
measure_model(tf_model, "Tensorflow", tf_batch)
nn.LeakyReLU(0.2, inplace=True)
])
self.block4 = nn.Sequential(*[
nn.Conv2d(256, 512, 4, padding=1),
nn.InstanceNorm2d(512, affine=True),
nn.LeakyReLU(0.2, inplace=True)
])
# FCN classification layer
self.fc = nn.Conv2d(512, 1, 4, padding=1)
def forward(self, x):
h1 = self.block1(x)
h2 = self.block2(h1)
h3 = self.block3(h2)
h4 = self.block4(h3)
out = self.fc(h4)
# Average pooling and flatten
out = F.avg_pool2d(out, out.size()[2:]).view(out.size()[0], -1)
D_outputs = namedtuple("DOutputs", ['h1', 'h2', 'h3', 'h4', 'out'])
return D_outputs(h1, h2, h3, h4, out)
if __name__ == "__main__":
from utils import measure_model, model_size
g = Generator(3, 3, transconv=False)
measure_model(g, 256, 256)
print(model_size(g))
def main(**kwargs):
global args, best_prec1
# Override if needed
for arg, v in kwargs.items():
args.__setattr__(arg, v)
### Calculate FLOPs & Param
model = getattr(models, args.model)(args)
if args.data in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
else:
IMAGE_SIZE = 224
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE,
args.debug)
if 'measure_only' in args and args.measure_only:
return n_flops, n_params
print('Starting.. FLOPs: %.2fM, Params: %.2fM' %
(n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % \
(args.model, int(n_params), int(n_flops))
del (model)
# Create model
model = getattr(models, args.model)(args)
if args.debug:
print(args)
print(model)
model = torch.nn.DataParallel(model).cuda()
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
# Resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
# Evaluate from a model
if args.evaluate_from is not None:
args.evaluate = True
state_dict = torch.load(args.evaluate_from)['state_dict']
model.load_state_dict(state_dict)
cudnn.benchmark = True
### Data loading
if args.data == "cifar10":
train_set = datasets.CIFAR10('../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR10('../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.data == "cifar100":
train_set = datasets.CIFAR100('../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR100('../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Run Forward / Backward passes
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
# Train for one epoch
tr_prec1, tr_prec5, loss, lr = \
train(train_loader, model, criterion, optimizer, epoch)
# Evaluate on validation set
val_prec1, val_prec5 = validate(val_loader, model, criterion)
# Remember best prec@1 and save checkpoint
is_best = val_prec1 < best_prec1
best_prec1 = max(val_prec1, best_prec1)
model_filename = 'checkpoint_%03d.pth.tar' % epoch
save_checkpoint(
{
'epoch': epoch,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, args, is_best, model_filename,
"%.4f %.4f %.4f %.4f %.4f %.4f\n" %
(val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr))
# TestModel and return
model = model.cpu().module
model = nn.DataParallel(model).cuda()
print(model)
validate(val_loader, model, criterion)
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('Finished training! FLOPs: %.2fM, Params: %.2fM' %
(n_flops / 1e6, n_params / 1e6))
print('Please run again with --resume --evaluate flags,'
' to evaluate the best model.')
return
def main(**kwargs):
global args, best_prec1
mapper = {
1 : 'airplane',
2 : 'automobile',
3 : 'bird',
4 : 'cat',
5 : 'deer',
6 : 'dog',
7 : 'frog',
8 : 'horse',
9 : 'ship',
10 : 'truck'
}
# Override if needed
for arg, v in kwargs.items():
args.__setattr__(arg, v)
### Calculate FLOPs & Param
model = getattr(models, args.model)(args)
if args.data in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
else:
IMAGE_SIZE = 224
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, args.debug)
if 'measure_only' in args and args.measure_only:
return n_flops, n_params
print('Starting.. FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % \
(args.model, int(n_params), int(n_flops))
del(model)
# Create model
model = getattr(models, args.model)(args)
if args.debug:
print(args)
print(model)
criterion = nn.CrossEntropyLoss()
if torch.cuda.is_available():
model = torch.nn.DataParallel(model).cuda()
criterion = criterion.cuda()
# Define loss function (criterion) and optimizer
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
checkpoint = load_checkpoint(args)
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
cudnn.benchmark = True
### Data loading
if args.data == "cifar10":
train_set = datasets.CIFAR10('../data', train=True, download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR10('../data', train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.data == "cifar100":
train_set = datasets.CIFAR100('../data', train=True, download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR100('../data', train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
val_loader = torch.utils.data.DataLoader(
val_set,
batch_size=2, shuffle=False,
num_workers=args.workers, pin_memory=True)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
try:
top1_per_cls = [AverageMeter() for i in range(0, model.module.num_blocks)]
top5_per_cls = [AverageMeter() for i in range(0, model.module.num_blocks)]
except:
top1_per_cls = [AverageMeter() for i in range(0, model.num_blocks)]
top5_per_cls = [AverageMeter() for i in range(0, model.num_blocks)]
model.eval()
end = time.time()
from gradCam import *
for i, (input, target) in enumerate(val_loader):
if args.imgNo > 0 and i!=args.imgNo:
continue
(inputM, targetM) = (input, target)
input, target = input[0], target[0].view(1)
input = input.view(1,input.shape[0], input.shape[1],input.shape[2])
if torch.cuda.is_available():
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input, requires_grad=True)
target_var = torch.autograd.Variable(target)
# ### Compute output
scores, feats = model(input_var, 0.0, p=1)
if args.model == 'msdnet':
loss = msd_loss(scores, target_var, criterion)
else:
loss = criterion(scores, target_var)
#############################################################################
name = 'diags/' + mapper[target_var.cpu().data.numpy()[0]+1] + '_' + str(args.imgNo) + '_' + str(args.maxC)
name2 = 'diags/' + mapper[target_var.cpu().data.numpy()[0]+1] + '_' + str(args.imgNo)
if len(args.classLabel):
if mapper[target_var.cpu().data.numpy()[0]+1] == args.classLabel:
print(correctNess(scores, target_var), i, mapper[target_var.cpu().data.numpy()[0]+1])
elif args.imgNo < 0:
print(correctNess(scores, target_var), i, mapper[target_var.cpu().data.numpy()[0]+1])
elif i == args.imgNo:
import os
if not os.path.exists('diags'):
os.makedirs('diags')
print("Category : {}".format(mapper[target_var.cpu().data.numpy()[0]+1]))
print(correctNess(scores, target_var))
grad_cam = GradCam(model = model)
mask = grad_cam(target_var.cpu().data.numpy()[0], input_var, 0, args.maxC-1)#target_index)
gb_model = GuidedBackpropReLUModel(model)
gb = gb_model(target_var.cpu().data.numpy()[0], input_var, 0, args.maxC-1).transpose(2,0,1)
img = input[0].cpu().data.numpy().transpose(1,2,0)
img = cv2.resize(img, (512, 512))
show_cam_on_image(img, mask, name)
utils.save_image(torch.from_numpy(gb), name +'_gb.jpg')
cam_mask = np.zeros(gb.shape)
for i in range(0, gb.shape[0]):
cam_mask[i, :, :] = mask
cam_gb = np.multiply(cam_mask, gb)
utils.save_image(torch.from_numpy(cam_gb), name +'_cam_gb.jpg')
img = cv2.resize(input.cpu().data.numpy()[0].transpose(1,2,0), (512, 512))
utils.save_image(torch.from_numpy(img.transpose(2,0,1)), name2 + '_input.jpg')
exit()
else:
continue
def main():
global args, model
### Calculate FLOPs & Param
model = getattr(models, args.model)(args)
IMAGE_SIZE = args.inp_img_size
n_flops, n_params = measure_model(model.cuda(), IMAGE_SIZE, IMAGE_SIZE,
args.scaling_factor)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % \
(args.model, int(n_params), int(n_flops))
del (model)
########### Building Model ##############
cudnn.benchmark = True
model = getattr(models, args.model)(args)
model = torch.nn.DataParallel(model).cuda()
criterion = charbonnier_loss().cuda()
optimizer = torch.optim.Adam(model.parameters(), args.lr)
lr = args.lr
########## Data Loading ###############
traindir = os.path.join("../data/", args.train_data)
testdir = os.path.join("../data/", args.test_data)
train_img = read_data(traindir,
crop_size=args.inp_img_size,
upscale_factor=args.scaling_factor,
c_dim=args.c_dim,
stride=128)
train_set = TensorClass(train_img[0], train_img[1], train_img[2])
test_set = testDatasetFromFolder2(testdir,
upscale_factor=args.scaling_factor)
train_loader = data.DataLoader(dataset=train_set,
num_workers=args.workers,
batch_size=args.batch_size,
shuffle=True)
test_loader = data.DataLoader(dataset=test_set,
num_workers=args.workers,
batch_size=test_batch_size,
shuffle=False)
### Optionally resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
### Optionally evaluate from a model
if args.evaluate_from is not None:
args.evaluate = True
state_dict = torch.load(args.evaluate_from)['state_dict']
model.load_state_dict(state_dict)
if args.evaluate:
test_psnr, test_ssim = test(test_loader, model, criterion)
print("test Psnr: " + str(test_psnr) + "\t test SSIM: " +
str(test_ssim))
return
for epoch in range(args.start_epoch, args.epochs):
### Train for one epoch
tr_psnr, tr_ssim ,loss, lr = \
train(train_loader, model, criterion, optimizer, epoch)
### Evaluate on test set
test_psnr, test_ssim = test(test_loader, model, criterion)
print("Train Psnr: " + str(tr_psnr) + "\t Train SSIM: " +
str(tr_ssim) + "\t Loss: " + str(loss))
print("test Psnr: " + str(test_psnr) + "\t test SSIM: " +
str(test_ssim))
### save checkpoint
model_filename = 'checkpoint_%03d.pth.tar' % epoch
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'psnr': best_psnr,
}, args, model_filename)
def main():
global args, best_acc1
args.filename = 'log.txt'
if args.dataset in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
else:
IMAGE_SIZE = 224
fd = open(args.record_file, 'a')
print('Args Config:', str(args))
fd.write(str(args) + '\n')
if args.continue_train:
print("Continue training!")
fd.write("Continue training!" + "\n")
args.resume = args.train_url + "save_models/checkpoint.pth.tar"
### Calculate FLOPs & Param
else:
if args.model.startswith("G"):
model = GhostNet.get_GhostNet(args)
else:
model = MicroNet.get_MicroNet(args)
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n')
print('Model Struture:', str(model))
fd.write(str(model) + '\n')
del (model)
fd.close()
# model = MicroNet.get_MicroNet(args)
if args.model.startswith("G"):
model = GhostNet.get_GhostNet(args)
else:
model = MicroNet.get_MicroNet(args)
model.cuda()
### Define loss function (criterion) and optimizer
if args.LSR_Mixup:
criterion = CrossEntryLoss_onehot().cuda()
else:
criterion = nn.CrossEntropyLoss().cuda()
criterion_temper = nn.CrossEntropyLoss().cuda() # for validation
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
### Optionally resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
cudnn.benchmark = True
if args.train_on_cloud:
args.data_url = "./data"
### Data loading
if args.dataset == "cifar10":
normalize = transforms.Normalize(mean=[0.4914, 0.4824, 0.4467],
std=[0.2471, 0.2435, 0.2616])
train_set = datasets.CIFAR10(args.data_url, train=True, download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4, padding_mode=args.pad_mode),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.CIFAR10(args.data_url, train=False,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]))
elif args.dataset == "cifar100":
normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408],
std=[0.2675, 0.2565, 0.2761])
train_set = datasets.CIFAR100(args.data_url, train=True, download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4, padding_mode=args.pad_mode),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.CIFAR100(args.data_url, train=False,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]))
elif args.dataset == "mini-imagenet":
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = mini_imagenet(args.data_url, "split_train.csv", transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = mini_imagenet(args.data_url, "split_val.csv", transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
else:
traindir = os.path.join(args.data_url, 'train')
valdir = os.path.join(args.data_url, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = datasets.ImageFolder(traindir, transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.ImageFolder(valdir, transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=False)
val_loader = torch.utils.data.DataLoader(
val_set,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=False)
if args.evaluate:
criterion_temper = nn.CrossEntropyLoss.cuda()
validate(val_loader, model,
criterion_temper) # TODO: validate must before measure_model otherwise result be worese, but I don't know why
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
fd = open(args.record_file, 'a')
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n')
return
epoch_time = AverageMeter()
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
### Train for one epoch
tr_acc1, tr_acc5, tr_loss, lr = \
train(train_loader, model, criterion, optimizer, epoch, args)
### Evaluate on validation set
val_acc1, val_acc5, val_loss = validate(val_loader, model, criterion_temper)
### Remember best Acc@1 and save checkpoint
is_best = val_acc1 > best_acc1
best_acc1 = max(val_acc1, best_acc1)
# model_filename = 'checkpoint_%03d.pth.tar' % epoch
model_filename = 'checkpoint.pth.tar'
save_checkpoint({
'epoch': epoch,
'model': args.model,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, args, is_best, model_filename, "%3d %.4f %.4f %.4f %.4f %.4f %.4f %.4f\n" %
(epoch, val_acc1, val_acc5, tr_acc1, tr_acc5, val_loss, tr_loss, lr))
epoch_time.update(time.time() - start_time, 1)
string = 'Avg_time: %4f H, Duration: %4f H, Left Time: %4f H' % \
(epoch_time.avg / 3600, epoch_time.sum / 3600, epoch_time.avg * (args.epochs - epoch - 1) / 3600)
print(string)
fd = open(args.record_file, 'a')
fd.write(string + '\n')
fd.close()
start_time = time.time()
### Convert model and test
# model = model.cpu().module
# print('Start Converting ...')
# convert_model(model, args)
# print('Converting End!')
print('Model Struture:', str(model))
validate(val_loader, model, criterion_temper)
model.cpu()
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
fd = open(args.record_file, 'a')
fd.write(str(model) + '\n')
fd.write('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6) + '\n')
fd.close()
return
def main():
global args, best_prec1
### Calculate FLOPs & Param
model = getattr(models, args.model)(args)
print(model)
if args.data in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
else:
IMAGE_SIZE = 224
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
del (model)
args.no_save_model = True
print(args)
### Create model
model = getattr(models, args.model)(args)
if args.model.startswith('alexnet') or args.model.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
### Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
### Optionally resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
### Optionally convert from a model and saves it to onnx file
if args.convert_from is not None:
args.evaluate = True
state_dict = torch.load(args.convert_from)['state_dict']
model.load_state_dict(state_dict)
model = model.cpu().module
convert_model(model, args)
model = model.cuda()
# head, tail = os.path.split(args.convert_from)
# tail = "converted_" + tail
# torch.save({'state_dict': model.state_dict()}, os.path.join(head, tail))
dummy_input = torch.randn(args.batch_size,
3,
IMAGE_SIZE,
IMAGE_SIZE,
device='cuda')
torch.onnx.export(model,
dummy_input,
args.convert_from + ".onnx",
verbose=True)
return
### Optionally evaluate from a model
if args.evaluate_from is not None:
args.evaluate = True
state_dict = torch.load(args.evaluate_from)['state_dict']
model.load_state_dict(state_dict)
cudnn.benchmark = True
### Data loading
if args.data == "cifar10":
norm_mean = [0.49139968, 0.48215827, 0.44653124]
norm_std = [0.24703233, 0.24348505, 0.26158768]
norm_transform = transforms.Normalize(norm_mean, norm_std)
train_set = datasets.CIFAR10('~/Documents/CIFAR-10',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(
32,
padding=4,
padding_mode='reflect'),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
norm_transform,
]))
val_set = datasets.CIFAR10('~/Documents/CIFAR-10',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
norm_transform,
]))
elif args.data == "cifar100":
normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408],
std=[0.2675, 0.2565, 0.2761])
train_set = datasets.CIFAR100('../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.CIFAR100('../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]))
else:
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
### Train for one epoch
tr_prec1, tr_prec5, loss, lr = \
train(train_loader, model, criterion, optimizer, epoch)
### Evaluate on validation set
val_prec1, val_prec5 = validate(val_loader, model, criterion)
### Remember best prec@1
is_best = val_prec1 < best_prec1
best_prec1 = max(val_prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, args, is_best, args.file_name)
### Convert model and test
model = model.cpu().module
convert_model(model, args)
model = nn.DataParallel(model).cuda()
print(model)
validate(val_loader, model, criterion)
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
return
def main(**kwargs):
global args, best_prec1, acc_file_name, flops_file_name
# Override if needed
# kwargs is a dict
for arg, v in kwargs.items():
args.__setattr__(arg, v) # set args.arg=v, copy kwargs to args
### Calculate FLOPs & Param
model = getattr(models, args.model)(args) ###### ? get models.(args.model)
if args.data in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
if (args.data == 'cifar10'):
acc_file_name = "cifar10_acc.csv"
flops_file_name = "cifar10_flops.csv"
else:
acc_file_name = "cifar100_acc.csv"
flops_file_name = "cifar100_flops.csv"
else:
IMAGE_SIZE = 224
set_save_path()
# calculate the FLOPs
# if (args.evaluate):
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE,
flops_file_path, args.debug) #####
if 'measure_only' in args and args.measure_only: # no 'measure_only' parameter
return n_flops, n_params
print('Starting.. FLOPs: %.2fM, Params: %.2fM' %
(n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % (args.model, int(n_params), int(n_flops))
# del(model)
# Create model(the 2nd)
# model = getattr(models, args.model)(args) ####
if args.debug:
print(args)
# print(model)
model = torch.nn.DataParallel(
model).cuda() # Implements data parallelism at the module level
# if exist mulit-devices, This container parallelizes the application of the given module by splitting
# the input across the specified devices by chunking in the batch dimension (other objects will be copied once per device)
### Data loading, no nomarlisation
if args.data == "cifar10":
train_set = datasets.CIFAR10('../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR10('../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.data == "cifar100":
train_set = datasets.CIFAR100('../data',
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
val_set = datasets.CIFAR100('../data',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
]))
cudnn.benchmark = True ######
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# test_loader = torch.utils.data.DataLoader(
# val_set,
# batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=True)
test_loader = val_loader
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
##### validate
# Resume from a checkpoint
if args.resume: # use latest checkpoint if have any
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
# Evaluate from a model
if args.evaluate_from is not None: # path to saved checkpoint
args.evaluate = True
state_dict = torch.load(args.evaluate_from)['state_dict']
model.load_state_dict(state_dict)
if args.evalmode is not None:
if args.evalmode == 'anytime':
validate(val_loader, model, criterion)
else:
print('dynamic')
dynamic_evaluate(model, test_loader, val_loader, args)
return
# Run Forward / Backward passes
# evaluate and return
# if args.evaluate:
# validate(val_loader, model, criterion)
# return
###### train
best_epoch = 0
for epoch in range(args.start_epoch, args.epochs):
# Train for one epoch
tr_prec1, tr_prec5, loss, lr = train(train_loader, model, criterion,
optimizer, epoch)
# Evaluate on validation set
val_prec1, val_prec5 = validate(val_loader, model, criterion)
# Remember best prec@1 and save checkpoint
is_best = val_prec1 < best_prec1
best_prec1 = max(val_prec1, best_prec1)
if (is_best):
best_epoch = epoch
model_filename = 'checkpoint_%03d.pth.tar' % epoch
save_checkpoint(
{
'epoch': epoch,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, args, is_best, model_filename,
"%.4f %.4f %.4f %.4f %.4f %.4f\n" %
(val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr))
print('Best val_prec1: {:.4f} at epoch {}'.format(best_prec1, best_epoch))
# TestModel and return
model = model.cpu().module
model = nn.DataParallel(model).cuda()
# if args.debug:
# print(model)
validate(val_loader, model, criterion)
# n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE, flops_file_path,args.debug)
# print('Finished training! FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
print('Please run again with --resume --evaluate flags,'
' to evaluate the best model.')
return
args = parser.parse_args()
# ------------------ MEAN & STD ---------------------
PIXEL_MEANS = np.array([0.485, 0.456, 0.406])
PIXEL_STDS = np.array([0.229, 0.224, 0.225])
# ---------------------------------------------------
# Set GPU id, CUDA and cudnn
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
USE_CUDA = torch.cuda.is_available()
cudnn.benchmark = True
# Create & Load model
MODEL = models.__dict__[args.arch]()
# Calculate FLOPs & Param
n_flops, n_convops, n_params = measure_model(MODEL, args.crop_size, args.crop_size)
print('==> FLOPs: {:.4f}M, Conv_FLOPs: {:.4f}M, Params: {:.4f}M'.
format(n_flops / 1e6, n_convops / 1e6, n_params / 1e6))
del MODEL
# Load Weights
MODEL = models.__dict__[args.arch]()
checkpoint = torch.load(args.model_weights)
weight_dict = checkpoint
model_dict = MODEL.state_dict()
updated_dict, match_layers, mismatch_layers = weight_filler(weight_dict, model_dict)
model_dict.update(updated_dict)
MODEL.load_state_dict(model_dict)
# Switch to evaluate mode
MODEL.cuda().eval()
def main():
global args, best_prec1
### Calculate FLOPs & Param
model = getattr(models, args.model)(args)
print(model)
if args.data in ['cifar10', 'cifar100']:
IMAGE_SIZE = 32
else:
IMAGE_SIZE = 224
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % \
(args.model, int(n_params), int(n_flops))
del (model)
print(args)
torch.manual_seed(args.manual_seed)
# torch.cuda.manual_seed_all(args.manual_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(args.manual_seed)
### Create model
model = getattr(models, args.model)(args)
if args.model.startswith('alexnet') or args.model.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
### Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
### Optionally resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
### Optionally convert from a model
if args.convert_from is not None:
args.evaluate = True
state_dict = torch.load(args.convert_from)['state_dict']
model.load_state_dict(state_dict)
model = model.cpu().module
convert_model(model, args)
model = nn.DataParallel(model).cuda()
head, tail = os.path.split(args.convert_from)
tail = "converted_" + tail
torch.save({'state_dict': model.state_dict()},
os.path.join(head, tail))
### Optionally evaluate from a model
if args.evaluate_from is not None:
args.evaluate = True
state_dict = torch.load(args.evaluate_from)['state_dict']
model.load_state_dict(state_dict)
### Data loading
dataset_dir = args.dataset_dir if args.dataset_dir is not None else '../data'
if args.data == "cifar10":
normalize = transforms.Normalize(mean=[0.4914, 0.4824, 0.4467],
std=[0.2471, 0.2435, 0.2616])
train_set = datasets.CIFAR10(dataset_dir,
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.CIFAR10(dataset_dir,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]))
elif args.data == "cifar100":
normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408],
std=[0.2675, 0.2565, 0.2761])
train_set = datasets.CIFAR100(dataset_dir,
train=True,
download=True,
transform=transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.CIFAR100(dataset_dir,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
normalize,
]))
else:
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_set = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomSizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_set = datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
### Train for one epoch
tr_prec1, tr_prec5, loss, lr = \
train(train_loader, model, criterion, optimizer, epoch)
### Evaluate on validation set
val_prec1, val_prec5 = validate(val_loader, model, criterion)
### Remember best prec@1 and save checkpoint
is_best = val_prec1 < best_prec1
best_prec1 = max(val_prec1, best_prec1)
model_filename = 'checkpoint_%03d.pth.tar' % epoch
save_checkpoint(
{
'epoch': epoch,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, args, is_best, model_filename,
"%.4f %.4f %.4f %.4f %.4f %.4f\n" %
(val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr))
### Convert model and test
model = model.cpu().module
convert_model(model, args)
model = nn.DataParallel(model).cuda()
print(model)
validate(val_loader, model, criterion)
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
return
def run_training(args):
# create model
model = models.__dict__[args.arch](args.pretrained)
model = torch.nn.DataParallel(model).cuda()
# measure FLOPs:
measure_model(model, 32, 32)
best_prec1 = 0
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
logging.info('=> loading checkpoint `{}`'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_iter = checkpoint['iter']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
logging.info('=> loaded checkpoint `{}` (iter: {})'.format(
args.resume, checkpoint['iter']
))
else:
logging.info('=> no checkpoint found at `{}`'.format(args.resume))
cudnn.benchmark = True
train_loader = prepare_train_data(dataset=args.dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
test_loader = prepare_test_data(dataset=args.dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
end = time.time()
if args.defend_algo == 'ifgm':
defend_algo = ifgm
elif args.defend_algo == 'None':
defend_algo = None
else:
raise NotImplementedError
# adversarial algorithm
if args.attack_algo == 'ifgm':
attack_algo = ifgm
elif args.attack_algo == 'None':
attack_algo = None
else:
raise NotImplementedError
for i in range(args.start_iter, args.iters):
model.train()
adjust_learning_rate(args, optimizer, i)
input, target = next(iter(train_loader))
# measuring data loading time
data_time.update(time.time() - end)
target = target.squeeze().long().cuda(async=True)
target_var = Variable(target)
if defend_algo:
input_adv = defend_algo(input, None, F.cross_entropy,
y=target_var,
eps=args.defend_eps,
model=model,
steps=args.defend_adv_iter,
gamma=args.defend_gamma,
randinit=args.defend_randinit).data
input_adv_var = Variable(input_adv)
input_var = Variable(input)
# compute output
output = model(input_var)
if defend_algo:
output_adv = model(input_adv_var)
loss = criterion(output, target_var) * 0.5 + criterion(output_adv, target_var) * 0.5
else:
loss = criterion(output, target_var)
# measure accuracy and record loss
prec1, = accuracy(output.data, target, topk=(1,))
losses.update(loss.item(), input.size(0))
top1.update(prec1, input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# print log
if i % args.print_freq == 0:
logging.info("Iter: [{0}/{1}]\t"
"Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
"Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
"Loss {loss.val:.3f} ({loss.avg:.3f})\t"
"Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t".format(
i,
args.iters,
batch_time=batch_time,
data_time=data_time,
loss=losses,
top1=top1)
)
# evaluate every 1000 steps
if (i % args.eval_every == 0 and i > 0) or (i == args.iters - 1):
prec1 = validate(args, test_loader, model, criterion)
if attack_algo:
_ = validate_adv(args, test_loader, model, criterion, attack_algo)
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
checkpoint_path = os.path.join(args.save_path,
'checkpoint_{:05d}.pth.tar'.format(
i))
save_checkpoint({
'iter': i,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
},
is_best, filename=checkpoint_path)
shutil.copyfile(checkpoint_path, os.path.join(args.save_path,
'checkpoint_latest'
'.pth.tar'))
def main():
global BEST_ACC, LR_STATE
start_epoch = cfg.CLS.start_epoch # start from epoch 0 or last checkpoint epoch
# Create ckpt folder
if not os.path.isdir(cfg.CLS.ckpt):
mkdir_p(cfg.CLS.ckpt)
if args.cfg_file is not None and not cfg.CLS.evaluate:
shutil.copyfile(
args.cfg_file,
os.path.join(cfg.CLS.ckpt,
args.cfg_file.split('/')[-1]))
# Dataset and Loader
normalize = transforms.Normalize(mean=cfg.pixel_mean, std=cfg.pixel_std)
train_aug = [
transforms.RandomResizedCrop(cfg.CLS.crop_size),
transforms.RandomHorizontalFlip()
]
if len(cfg.CLS.rotation) > 0:
train_aug.append(transforms.RandomRotation(cfg.CLS.rotation))
if len(cfg.CLS.pixel_jitter) > 0:
train_aug.append(RandomPixelJitter(cfg.CLS.pixel_jitter))
if cfg.CLS.grayscale > 0:
train_aug.append(transforms.RandomGrayscale(cfg.CLS.grayscale))
train_aug.append(transforms.ToTensor())
train_aug.append(normalize)
traindir = os.path.join(cfg.CLS.data_root, cfg.CLS.train_folder)
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir, transforms.Compose(train_aug)),
batch_size=cfg.CLS.train_batch,
shuffle=True,
num_workers=cfg.workers,
pin_memory=True)
if cfg.CLS.validate or cfg.CLS.evaluate:
valdir = os.path.join(cfg.CLS.data_root, cfg.CLS.val_folder)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(cfg.CLS.base_size),
transforms.CenterCrop(cfg.CLS.crop_size),
transforms.ToTensor(),
normalize,
])),
batch_size=cfg.CLS.test_batch,
shuffle=False,
num_workers=cfg.workers,
pin_memory=True)
# Create model
model = models.__dict__[cfg.CLS.arch]()
print(model)
# Calculate FLOPs & Param
n_flops, n_convops, n_params = measure_model(model, cfg.CLS.crop_size,
cfg.CLS.crop_size)
print('==> FLOPs: {:.4f}M, Conv_FLOPs: {:.4f}M, Params: {:.4f}M'.format(
n_flops / 1e6, n_convops / 1e6, n_params / 1e6))
del model
model = models.__dict__[cfg.CLS.arch]()
# Load pre-train model
if cfg.CLS.pretrained:
print("==> Using pre-trained model '{}'".format(cfg.CLS.pretrained))
pretrained_dict = torch.load(cfg.CLS.pretrained)
try:
pretrained_dict = pretrained_dict['state_dict']
except:
pretrained_dict = pretrained_dict
model_dict = model.state_dict()
updated_dict, match_layers, mismatch_layers = weight_filler(
pretrained_dict, model_dict)
model_dict.update(updated_dict)
model.load_state_dict(model_dict)
else:
print("==> Creating model '{}'".format(cfg.CLS.arch))
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
if cfg.CLS.pretrained:
def param_filter(param):
return param[1]
new_params = map(
param_filter,
filter(lambda p: p[0] in mismatch_layers,
model.named_parameters()))
base_params = map(
param_filter,
filter(lambda p: p[0] in match_layers, model.named_parameters()))
model_params = [{
'params': base_params
}, {
'params': new_params,
'lr': cfg.CLS.base_lr * 10
}]
else:
model_params = model.parameters()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
optimizer = optim.SGD(model_params,
lr=cfg.CLS.base_lr,
momentum=cfg.CLS.momentum,
weight_decay=cfg.CLS.weight_decay)
# Evaluate model
if cfg.CLS.evaluate:
print('\n==> Evaluation only')
test_loss, test_top1, test_top5 = test(val_loader, model, criterion,
start_epoch, USE_CUDA)
print(
'==> Test Loss: {:.8f} | Test_top1: {:.4f}% | Test_top5: {:.4f}%'.
format(test_loss, test_top1, test_top5))
return
# Resume training
title = 'Pytorch-CLS-' + cfg.CLS.arch
if cfg.CLS.resume:
# Load checkpoint.
print("==> Resuming from checkpoint '{}'".format(cfg.CLS.resume))
assert os.path.isfile(
cfg.CLS.resume), 'Error: no checkpoint directory found!'
checkpoint = torch.load(cfg.CLS.resume)
BEST_ACC = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
# Train and val
for epoch in range(start_epoch, cfg.CLS.epochs):
print('\nEpoch: [{}/{}] | LR: {:.8f}'.format(epoch + 1, cfg.CLS.epochs,
LR_STATE))
train_loss, train_acc = mixup_train(train_loader, model, criterion,
optimizer, epoch, USE_CUDA)
if cfg.CLS.validate:
test_loss, test_top1, test_top5 = test(val_loader, model,
criterion, epoch, USE_CUDA)
else:
test_loss, test_top1, test_top5 = 0.0, 0.0, 0.0
# Append logger file
logger.append([LR_STATE, train_loss, test_loss, train_acc, test_top1])
# Save model
save_checkpoint(model, optimizer, test_top1, epoch)
# Draw curve
try:
draw_curve(cfg.CLS.arch, cfg.CLS.ckpt)
print('==> Success saving log curve...')
except:
print('==> Saving log curve error...')
logger.close()
try:
savefig(os.path.join(cfg.CLS.ckpt, 'log.eps'))
shutil.copyfile(
os.path.join(cfg.CLS.ckpt, 'log.txt'),
os.path.join(
cfg.CLS.ckpt, 'log{}.txt'.format(
datetime.datetime.now().strftime('%Y%m%d%H%M%S'))))
except:
print('copy log error.')
print('==> Training Done!')
print('==> Best acc: {:.4f}%'.format(BEST_ACC))
def main():
global args
### Create model and calculate FLOPs & Param and delete model
model = ShuffleNetV2(args.symbol)
print(model)
IMAGE_SIZE = 224
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
args.filename = "%s_%s_%s.txt" % (args.model, int(n_params), int(n_flops))
print(args)
del (model)
### Create model again
model = ShuffleNetV2(args.symbol)
# print(model)
# args.filename = "%s.txt" % (args.model)
# print(args)
model = torch.nn.DataParallel(model).cuda()
### Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
### Optionally resume from a checkpoint
if args.resume:
checkpoint = load_checkpoint(args)
if checkpoint is not None:
args.start_epoch = checkpoint['epoch'] + 1
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
best_prec1 = 0
cudnn.benchmark = True
### Data loading
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if not args.evaluate:
traindir = os.path.join(args.data, 'train')
train_set = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224,
scale=(0.2, 1.0),
ratio=(0.75, 1.33)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]))
valdir = os.path.join(args.data, 'val')
val_set = datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
]))
if not args.evaluate:
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
### Evaluate
if args.evaluate:
validate(val_loader, model, criterion)
return
### Train
for epoch in range(args.start_epoch, args.epochs):
### Train for one epoch
tr_prec1, tr_prec5, loss, lr = train(train_loader, model, criterion,
optimizer, epoch)
### Evaluate on validation set
val_prec1, val_prec5 = validate(val_loader, model, criterion)
### Remember best prec@1 and save checkpoint
is_best = (val_prec1 > best_prec1)
best_prec1 = max(val_prec1, best_prec1)
### Save checkpoint
model_filename = 'checkpoint_%03d.pth.tar' % epoch
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict()
}, args, is_best, model_filename,
"%.4f %.4f %.4f %.4f %.4f %.4f\n" %
(val_prec1, val_prec5, tr_prec1, tr_prec5, loss, lr))
# save_checkpoint({'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()},
# args, model_filename, "%.4f %.4f %.4f %.4f\n" % (tr_prec1, tr_prec5, loss, lr))
n_flops, n_params = measure_model(model, IMAGE_SIZE, IMAGE_SIZE)
print('FLOPs: %.2fM, Params: %.2fM' % (n_flops / 1e6, n_params / 1e6))
return
