def main(args):
flow.env.init()
flow.enable_eager_execution()
start_t = time.time()
res50_module = resnet50(num_classes=8)
end_t = time.time()
print("init time : {}".format(end_t - start_t))
start_t = time.time()
pretrain_models = flow.load(args.model_path)
res50_module.load_state_dict(pretrain_models)
end_t = time.time()
print("load params time : {}".format(end_t - start_t))
res50_module.eval()
res50_module.to("cuda")
start_t = time.time()
image = load_image(args.image_path)
image = flow.Tensor(image, device=flow.device("cuda"))
logits, body = res50_module(image)
predictions = logits.softmax()
predictions = predictions.numpy()
end_t = time.time()
print("infer time : {}".format(end_t - start_t))
clsidx = np.argmax(predictions)
print("predict prob: %f, class name: %s" %
(np.max(predictions), clsidx_2_labels[clsidx]))
def main(args):
start_t = time.perf_counter()
print("***** Model Init *****")
model = resnet50()
model.load_state_dict(flow.load(args.model_path))
model = model.to("cuda")
model.eval()
end_t = time.perf_counter()
print(f"***** Model Init Finish, time escapled {end_t - start_t:.6f} s *****")
if args.graph:
model_graph = InferGraph(model)
start_t = end_t
image = load_image(args.image_path)
image = flow.Tensor(image, device=flow.device("cuda"))
if args.graph:
pred = model_graph(image)
else:
pred = model(image).softmax()
pred = pred.numpy()
prob = np.max(pred)
clsidx = np.argmax(pred)
cls = clsidx_2_labels[clsidx]
end_t = time.perf_counter()
print(
"predict image ({}) prob: {:.5f}, class name: {}, time escapled: {:.6f} s".format(
os.path.basename(args.image_path), prob, cls, end_t - start_t
)
)
def __init__(self):
args = get_args()
for k, v in args.__dict__.items():
setattr(self, k, v)
self.rank = flow.env.get_rank()
self.world_size = flow.env.get_world_size()
self.cur_epoch = 0
self.cur_iter = 0
self.cur_batch = 0
self.is_global = (self.world_size > 1 and not self.ddp) or self.graph
self.is_train = False
self.meter_lr = self.graph is False
self.init_logger()
flow.boxing.nccl.set_fusion_threshold_mbytes(
self.nccl_fusion_threshold_mb)
flow.boxing.nccl.set_fusion_max_ops_num(self.nccl_fusion_max_ops)
if self.use_fp16 and self.num_nodes * self.num_devices_per_node > 1:
flow.boxing.nccl.enable_use_buffer_to_fuse_all_reduce(False)
self.model = resnet50(
zero_init_residual=self.zero_init_residual,
fuse_bn_relu=self.fuse_bn_relu,
fuse_bn_add_relu=self.fuse_bn_add_relu,
channel_last=self.channel_last,
)
self.init_model()
self.cross_entropy = make_cross_entropy(args)
self.train_data_loader = make_data_loader(args, "train",
self.is_global,
self.synthetic_data)
self.val_data_loader = make_data_loader(args, "validation",
self.is_global,
self.synthetic_data)
self.optimizer = make_optimizer(args, self.model)
self.lr_scheduler = make_lr_scheduler(args, self.optimizer)
self.acc = Accuracy()
if self.graph:
self.train_graph = make_train_graph(
self.model,
self.cross_entropy,
self.train_data_loader,
self.optimizer,
self.lr_scheduler,
return_pred_and_label=self.metric_train_acc,
)
self.eval_graph = make_eval_graph(self.model, self.val_data_loader)
if self.gpu_stat_file is not None:
self.gpu_stat = CudaUtilMemStat(f"rank{self.rank}_" +
self.gpu_stat_file,
only_ordinal=self.rank)
else:
self.gpu_stat = None
def get_model(dataset, model_path, is_cnn):
if dataset == 'cifar10':
from models import resnet20
model = resnet20.resnet20()
model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
model.load_state_dict(torch.load(model_path))
model = model.module
# weight_clip(model, adder2d, [-10, 10])
# low_pricision(model, adder2d, 4)
elif dataset == 'cifar10_cnn':
from models.resnet_cnn import cifar_resnet20
model = cifar_resnet20(pretrained='cifar10')
# Do quantilization
# model.eval()
# model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
# torch.quantization.fuse_modules(model, [['conv1', 'bn1']], inplace=True)
# for module in model.modules():
# if isinstance(module, BasicBlock):
# torch.quantization.fuse_modules(module, [['conv1', 'bn1'], ['conv2', 'bn2']], inplace=True)
# for module in model.modules():
# if isinstance(module, torch.nn.Sequential) and len(module) == 2:
# torch.quantization.fuse_modules(module, ['0', '1'], inplace=True)
# print(model)
# model_fp32_prepared = torch.quantization.prepare(model)
# low_pricision(model, torch.nn.Conv2d, 6)
elif dataset == 'ImageNet':
from models import resnet50
model = resnet50.resnet50()
model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
model.load_state_dict(torch.load(model_path))
model = model.module
elif dataset == "PACS":
# if is_cnn:
# from torchvision.models import resnet50
# model = resnet50(num_classes=7)
# else:
# from models.resnet50 import resnet50
# model = resnet50(num_classes=7)
#
# model.load_state_dict(torch.load(model_path))
model = torch.load(model_path)
else:
assert False
return model
def __init__(self, ckpt_dir, params=None):
self.ckpt_dir = ckpt_dir
if not params:
params = self.load_default_params()
self.params = params
# Sess start
if not self.params['gpu']:
self.sess = tf.Session()
else:
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
self.sess = tf.Session(config=config)
self.image_input = tf.placeholder(tf.float32, [None, 112, 112, 3], name='image_inputs')
self.image_label_input = tf.placeholder(tf.int64, [None, ], name='labels_inputs')
with tf.name_scope('face_resnet50'):
self.emb = resnet50(self.image_input, is_training=False)
print(self.image_input)
print(self.image_label_input)
print(self.emb)
load_graph_from_ckpt(self.sess, ckpt_dir, graph=False)
def __init__(self,
num_classes=1,
num_channels=3,
pretrained=False,
pretrained_model_path=""):
super().__init__()
assert num_channels == 3
self.num_classes = num_classes
filters = [64, 128, 256, 512]
resnet = resnet50()
# if pretrained=True, load pretrained resnet model
if pretrained:
resnet.load_state_dict(flow.load(pretrained_model_path))
self.firstconv = resnet.conv1
self.firstbn = resnet.bn1
self.firstrelu = resnet.relu
self.firstmaxpool = resnet.maxpool
self.encoder1 = resnet.layer1
self.encoder2 = resnet.layer2
self.encoder3 = resnet.layer3
self.encoder4 = resnet.layer4
# # Decoder
self.decoder4 = DecoderBlockLinkNet(filters[3], filters[2])
self.decoder3 = DecoderBlockLinkNet(filters[2], filters[1])
self.decoder2 = DecoderBlockLinkNet(filters[1], filters[0])
self.decoder1 = DecoderBlockLinkNet(filters[0], filters[0])
# Final Classifier
self.finaldeconv1 = nn.ConvTranspose2d(filters[0], 32, 3, stride=2)
self.finalrelu1 = nn.ReLU(inplace=True)
self.finalconv2 = nn.Conv2d(32, 32, 3)
self.finalrelu2 = nn.ReLU(inplace=True)
self.finalconv3 = nn.Conv2d(32, num_classes, 2, padding=1)
stratify=parsed_args['stratify'])
trainset = datasets.get_trainset()
valset = datasets.get_valset()
print('Train set is of length', trainset.__len__())
print('Val set is of length', valset.__len__())
batch_size = parsed_args["batch_size"]
shuffle = parsed_args['shuffle']
#Create the batched loaders
trainloader = create_batched_loader(trainset,
batch_size=batch_size,
shuffle=shuffle)
valloader = create_batched_loader(valset,
batch_size=batch_size,
shuffle=shuffle)
#Define the model
model = resnet50(num_classes=loaded_dataset.num_classes())
#Define the loss_fn
loss_fn = nn.CrossEntropyLoss()
#Define the optimizer
lr = parsed_args['lr']
opt = optim.Adam(model.parameters(), lr=0.001)
#Define the scheduler
scheduler = optim.lr_scheduler.MultiStepLR(opt, [5, 10])
resume = False
if parsed_args['resume'] is not None:
resume = True
# Try weight loading, if weight file provided.
try:
weight_file = str(Path(parsed_args['weight_file'][0]).absolute())
model = load_init_weights(weight_file,
parsed_args['weight_type'][0],
def train(configurations):
"""
:param configurations: a dictionary
:return:
"""
# Load images and labels for trainset and
# Load images and labels for testset
train_images, train_labels = get_train(configurations)
test_images, test_labels = get_test(configurations)
# Change data type to numpy
train_images = np.asarray(train_images)
test_images = np.asarray(test_images)
# Preprocess all the train images
train_images, test_images = preproc(train_images, test_images)
# Make labels to one-hot vector type
train_labels = one_hot(train_labels, configurations)
test_labels = one_hot(test_labels, configurations)
# Batch
batch_train = Batch(train_images, train_labels,
configurations['batch_size'])
batch_test = Batch(test_images, test_labels, configurations['batch_size'])
# Define tensors and operations on tensorflow
x = tf.placeholder(tf.float32,
shape=[
configurations['batch_size'],
configurations['height'], configurations['width'],
configurations['num_channels']
],
name='x')
y = tf.placeholder(
tf.float32,
shape=[configurations['batch_size'], configurations['num_classes']],
name='y')
logits_op = resnet50(x)
probs_op = tf.nn.sigmoid(logits_op)
loss_op = sigmoid_cross_entropy_loss(labels=y, logits=logits_op)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
trainer_op = tf.train.AdamOptimizer(
learning_rate=configurations['learning_rate']).minimize(loss_op)
saver = tf.train.Saver()
with tf.Session() as sess:
lowest_loss = None
tf.global_variables_initializer().run()
for epoch in range(1, configurations['num_epochs'] + 1):
train_x, train_y = batch_train.next_to()
loss, prob, _ = sess.run([loss_op, probs_op, trainer_op],
feed_dict={
x: augment(train_x, [224, 224]),
y: train_y
})
acc = get_acc(prob, train_y, 0.05)
print("epoch: {} - loss: {} - acc: {}".format(epoch, loss, acc))
# Save a model that outputs a lowest loss until now
if lowest_loss == None or lowest_loss > loss:
print("lowest loss: {}".format(loss))
lowest_loss = loss
saver.save(
sess,
os.path.join(configurations['ckpt_path'], 'model.ckpt'))
if epoch % configurations['test_cycles'] == 0:
test_x, test_y = batch_test.next_to()
loss, prob = sess.run([loss_op, probs_op],
feed_dict={
x: resize(test_x, [224, 224]),
y: test_y
})
acc = get_acc(prob, test_y, 0.05)
print("In test phase loss: {} - acc: {}".format(loss, acc))
import torch
from torch import cuda
from torch import optim, nn
from models.resnet50 import resnet50
from torchvision import transforms
input_size=(256,256,3)
classes=['normal','monotone','screenshot','unknown']
config={
'batch_size':64,
'train_split':0.8,
'lr':0.000005,
'n_epochs':100,
'print_every':1,
'cuda':True if cuda.is_available() else False,
'model':resnet50(pretrained=True,freeze=True),
'criterion':nn.CrossEntropyLoss(),
'transform':transforms.Compose([
transforms.Resize(size=224),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
]),
'early_stop':100,
'balance':True
}
config['optimizer']=optim.Adam(config['model'].parameters(), lr=config['lr'])
config['scheduler']=optim.lr_scheduler.ReduceLROnPlateau(config['optimizer'])
def main_worker(gpu, args):
global best_acc1
global best_auc
global minimum_loss
global count
global best_accdr
args.gpu = gpu
if not os.path.exists(args.model_dir):
os.makedirs(args.model_dir)
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.arch == "vgg11":
from models.vgg import vgg11
model = vgg11(num_classes=args.num_class, crossCBAM=args.crossCBAM)
elif args.arch == "resnet50":
from models.resnet50 import resnet50
model = resnet50(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen, CAN_TS=args.CAN_TS, crossCBAM=args.crossCBAM,
crosspatialCBAM = args.crosspatialCBAM, choice=args.choice)
elif args.arch == "resnet34":
from models.resnet50 import resnet34
model = resnet34(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen,CAN_TS=args.CAN_TS, crossCBAM=args.crossCBAM,
crosspatialCBAM = args.crosspatialCBAM)
elif args.arch == "resnet18":
from models.resnet50 import resnet18
model = resnet18(num_classes=args.num_class, multitask=args.multitask, liu=args.liu,
chen=args.chen, flagCBAM=False, crossCBAM=args.crossCBAM)
elif args.arch == "densenet161":
from models.densenet import densenet161
model = densenet161(num_classes=args.num_class, multitask=args.multitask, cosface=False, liu=args.liu,
chen=args.chen, crossCBAM=args.crossCBAM)
elif args.arch == "wired":
from models.wirednetwork import CNN
model = CNN(args, num_classes=args.num_class)
else:
print ("no backbone model")
if args.pretrained:
print ("==> Load pretrained model")
model_dict = model.state_dict()
pretrain_path = {"resnet50": "pretrain/resnet50-19c8e357.pth",
"resnet34": "pretrain/resnet34-333f7ec4.pth",
"resnet18": "pretrain/resnet18-5c106cde.pth",
"densenet161": "pretrain/densenet161-8d451a50.pth",
"vgg11": "pretrain/vgg11-bbd30ac9.pth",
"densenet121": "pretrain/densenet121-a639ec97.pth"}[args.arch]
pretrained_dict = torch.load(pretrain_path)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
pretrained_dict.pop('classifier.weight', None)
pretrained_dict.pop('classifier.bias', None)
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0))
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
if args.adam:
optimizer = torch.optim.Adam(model.parameters(), args.base_lr, weight_decay=args.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(), args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location={'cuda:4':'cuda:0'})
# args.start_epoch = checkpoint['epoch']
# load partial weights
if not args.evaluate:
print ("load partial weights")
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in checkpoint['state_dict'].items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
else:
print("load whole weights")
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
exit(0)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
size = 224
tra = transforms.Compose([
transforms.Resize(256),
transforms.RandomResizedCrop(size),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
# transforms.RandomRotation(90),
# transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),
transforms.ToTensor(),
normalize,
])
tra_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize])
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
# transforms.RandomRotation([-180, 180]),
# transforms.RandomAffine([-180, 180], translate=[0.1, 0.1], scale=[0.7, 1.3]),
# transforms.RandomCrop(224),
# # transforms.CenterCrop(224),
# transforms.ToTensor(),
# normalize
# ])
# print (args.model_dir)
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
# transforms.RandomRotation([-180, 180]),
# transforms.RandomAffine([-180, 180], translate=[0.1, 0.1], scale=[0.7, 1.3]),
# transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
# transforms.ToTensor(),
# normalize
# ])
# tra_test = transforms.Compose([
# transforms.Resize(350),
# transforms.CenterCrop(224),
# transforms.ToTensor(),
# normalize])
if args.dataset == 'amd':
from datasets.amd_dataset import traindataset
elif args.dataset == 'pm':
from datasets.pm_dataset import traindataset
elif args.dataset == "drdme":
from datasets.drdme_dataset import traindataset
elif args.dataset == "missidor":
from datasets.missidor import traindataset
elif args.dataset == "kaggle":
from datasets.kaggle import traindataset
else:
print ("no dataset")
exit(0)
val_dataset = traindataset(root=args.data, mode = 'val',
transform=tra_test, num_class=args.num_class,
multitask=args.multitask, args=args)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.evaluate:
a = time.time()
# savedir = args.resume.replace("model_converge.pth.tar","")
savedir = args.resume.replace(args.resume.split("/")[-1], "")
# savedir = "./"
if not args.multitask:
acc, auc, precision_dr, recall_dr, f1score_dr = validate(val_loader, model, args)
result_list = [acc, auc, precision_dr, recall_dr, f1score_dr]
print ("acc, auc, precision, recall, f1", acc, auc, precision_dr, recall_dr, f1score_dr)
save_result_txt(savedir, result_list)
print("time", time.time() - a)
return
else:
acc_dr, acc_dme, acc_joint, other_results, se, sp = validate(val_loader, model, args)
print ("acc_dr, acc_dme, acc_joint", acc_dr, acc_dme, acc_joint)
exit(0)
print ("auc_dr, auc_dme, precision_dr, precision_dme, recall_dr, recall_dme, f1score_dr, f1score_dme",
other_results)
print ("se, sp", se, sp)
result_list = [acc_dr, acc_dme, acc_joint]
result_list += other_results
result_list += [se, sp]
save_result_txt(savedir, result_list)
print ("time", time.time()-a)
return
train_dataset = traindataset(root=args.data, mode='train', transform=tra, num_class=args.num_class,
multitask=args.multitask, args=args)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True,worker_init_fn=worker_init_fn)
writer = SummaryWriter()
writer.add_text('Text', str(args))
#
from lr_scheduler import LRScheduler
lr_scheduler = LRScheduler(optimizer, len(train_loader), args)
for epoch in range(args.start_epoch, args.epochs):
is_best = False
is_best_auc = False
is_best_acc = False
# lr = adjust_learning_rate(optimizer, epoch, args)
# writer.add_scalar("lr", lr, epoch)
# train for one epoch
loss_train = train(train_loader, model, criterion, lr_scheduler, writer, epoch, optimizer, args)
writer.add_scalar('Train loss', loss_train, epoch)
# evaluate on validation set
if epoch % 5 == 0:
if args.dataset == "kaggle":
acc_dr, auc_dr = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val auc_dr", auc_dr, epoch)
is_best = acc_dr >= best_acc1
best_acc1 = max(acc_dr, best_acc1)
elif not args.multitask:
acc, auc, precision, recall, f1 = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc, epoch)
writer.add_scalar("Val auc_dr", auc, epoch)
is_best = auc >= best_acc1
best_acc1 = max(auc, best_acc1)
else:
acc_dr, acc_dme, joint_acc, other_results, se, sp , losses = validate(val_loader, model, args,criterion)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val acc_dme", acc_dme, epoch)
writer.add_scalar("Val acc_joint", joint_acc, epoch)
writer.add_scalar("Val auc_dr", other_results[0], epoch)
writer.add_scalar("Val auc_dme", other_results[1], epoch)
writer.add_scalar("val loss", losses, epoch)
is_best = joint_acc >= best_acc1
best_acc1 = max(joint_acc, best_acc1)
is_best_auc = other_results[0] >= best_auc
best_auc = max(other_results[0], best_auc)
is_best_acc = acc_dr >= best_accdr
best_accdr = max(acc_dr, best_accdr)
if not args.invalid:
if is_best:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer' : optimizer.state_dict(),
}, is_best, filename = "model_converge.pth.tar", save_dir=args.model_dir)
if is_best_auc:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_auc,
'optimizer' : optimizer.state_dict(),
}, False, filename = "converge_auc.pth.tar", save_dir=args.model_dir)
if is_best_acc:
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_accdr,
'optimizer': optimizer.state_dict(),
}, False, filename="converge_acc.pth.tar", save_dir=args.model_dir)
##convert .pth model to onnx model
import torch
import torchvision.models as models
#
import os
import sys
# sys.path.append(os.getcwd())
sys.path.append("/home/zigangzhao/DMS/dms0715/dms-5/")
from models.resnet50 import resnet50
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch_model = torch.load(
"./resnet50-42-best-0.9927184581756592.pth") # pytorch模型加载
model = resnet50()
model.fc = torch.nn.Linear(2048, 5)
model.load_state_dict(torch_model)
batch_size = 1 #批处理大小
input_shape = (3, 224, 384) #输入数据
# # set the model to inference mode
#torch_model.eval()
model.eval()
x = torch.randn(batch_size, *input_shape) # 生成张量
# x = x.to(device)
export_onnx_file = "0727.onnx" # 目的ONNX文件名
torch.onnx.export(
model,
def train(self):
from models.resnet50 import resnet50
from losses.loss import arcface_loss
tf.reset_default_graph()
num_classes = self.params['num_classes']
batch_size = self.params['batch_size']
ckpt_save_dir = self.params['ckpt/']
weight_decay = self.params['weight_decay']
max_ep = self.params['max_ep']
w_init_method = tf.contrib.layers.xavier_initializer(uniform=False)
tfr = self.params['tfrecords']
dataset = tf.data.TFRecordDataset(tfr)
dataset = dataset.map(parse_function)
dataset = dataset.shuffle(buffer_size=10000)
dataset = dataset.batch(batch_size)
iterator = dataset.make_initializable_iterator()
next_element = iterator.get_next()
image_input = tf.placeholder(tf.float32, [None, 112, 112, 3], name='image_inputs')
image_label_input = tf.placeholder(tf.int64, [None, ], name='labels_inputs')
with tf.name_scope('face_resnet50'):
emb = resnet50(image_input, is_training=True)
logit = arcface_loss(embedding=emb, labels=labels, w_init=w_init_method, out_num=num_classes)
inference_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logit, labels=labels))
conv_var = [var for var in tf.trainable_variables() if 'conv' in var.name]
l2_loss = tf.add_n([tf.nn.l2_loss(var) for var in conv_var])
loss = inference_loss + weight_decay * l2_loss
global_step = tf.Variable(name='global_step', initial_value=0, trainable=False)
inc_op = tf.assign_add(global_step, 1, name='increment_global_step')
lr_schedule = tf.train.piecewise_constant(global_step, boundaries=self.params['lr_steps'], values=self.params['lr'], name='lr_schedule')
opt = tf.train.MomentumOptimizer(learning_rate=lr_schedule, momentum=0.9)
grads = opt.compute_gradients(loss)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = opt.apply_gradients(grads, global_step=global_step)
pred = tf.nn.softmax(logit)
acc = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, axis=1), labels), dtype=tf.float32))
saver = tf.train.Saver(max_to_keep=3)
config = tf.ConfigProto()
config.allow_soft_placement = True
config.gpu_options.allow_growth = True
counter = 0
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
for i in range(max_ep):
sess.run(iterator.initializer)
while True:
try:
image_train, label_train = sess.run(next_element)
feed_dict = {images: image_train, labels: label_train}
_, loss_val, acc_val, _ = sess.run([train_op, inference_loss, acc, inc_op], feed_dict=feed_dict)
counter += 1
# print('counter: ', counter, 'loss_val', loss_val, 'acc: ', acc_val)
if counter % 100 == 0:
print('counter: ', counter, 'loss_val', loss_val, 'acc: ', acc_val)
filename = 'Face_iter_{:d}'.format(counter) + '.ckpt'
filename = os.path.join(ckpt_save_dir, filename)
saver.save(sess, filename)
except tf.errors.OutOfRangeError:
print('End of epoch %d', i)
break
import torch
from torchsummary import summary
from models.resnet50 import resnet50
model = resnet50(sobel=True)
model.cuda()
summary(model, (3, 224, 224))
f1 = open('resnet_debug_arch.txt', 'w')
for name, param in model.state_dict().items():
l = name+'\t'+str(param.size())
f1.write(l+'\n')
f1.close()
def main_worker(gpu, args):
global best_acc1
global minimum_loss
global count
args.gpu = gpu
if not os.path.exists(args.model_dir):
os.makedirs(args.model_dir)
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.arch == "vgg11":
from models.vgg import vgg11
model = vgg11(num_classes=args.num_class, crossCBAM=args.crossCBAM)
elif args.arch == "resnet50":
from models.resnet50 import resnet50
model = resnet50(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "resnet34":
from models.resnet50 import resnet34
model = resnet34(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "resnet18":
from models.resnet50 import resnet18
model = resnet18(num_classes=args.num_class,
multitask=args.multitask,
liu=args.liu,
chen=args.chen,
flagCBAM=False,
crossCBAM=args.crossCBAM)
elif args.arch == "densenet161":
from models.densenet import densenet161
model = densenet161(num_classes=args.num_class,
multitask=args.multitask,
cosface=False,
liu=args.liu,
chen=args.chen)
elif args.arch == "wired":
from models.wirednetwork import CNN
model = CNN(args, num_classes=args.num_class)
else:
print("no backbone model")
if args.pretrained:
print("==> Load pretrained model")
model_dict = model.state_dict()
pretrain_path = {
"resnet50": "pretrain/resnet50-19c8e357.pth",
"resnet34": "pretrain/resnet34-333f7ec4.pth",
"resnet18": "pretrain/resnet18-5c106cde.pth",
"densenet161": "pretrain/densenet161-8d451a50.pth",
"vgg11": "pretrain/vgg11-bbd30ac9.pth"
}[args.arch]
pretrained_dict = torch.load(pretrain_path)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
pretrained_dict.pop('classifier.weight', None)
pretrained_dict.pop('classifier.bias', None)
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.weight_decay)
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr,
# momentum=args.momentum,
# weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume,
map_location={'cuda:4': 'cuda:0'})
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
exit(0)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
size = 224
# tra = transforms.Compose([
# # transforms.Resize(256),
# transforms.RandomResizedCrop(size),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# # transforms.RandomRotation(90),
# # transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),
# transforms.ToTensor(),
# normalize,
# ])
# tra_test = transforms.Compose([
# transforms.Resize(size+32),
# transforms.CenterCrop(size),
# transforms.ToTensor(),
# normalize])
tra = transforms.Compose([
transforms.Resize(350),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
# transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
transforms.RandomRotation([-180, 180]),
transforms.RandomAffine([-180, 180],
translate=[0.1, 0.1],
scale=[0.7, 1.3]),
transforms.RandomCrop(224),
# transforms.CenterCrop(224),
transforms.ToTensor(),
normalize
])
# tra = transforms.Compose([
# transforms.Resize(350),
# transforms.RandomHorizontalFlip(),
# transforms.RandomVerticalFlip(),
# transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
# transforms.ToTensor(),
# normalize])
#
tra_test = transforms.Compose([
transforms.Resize(350),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
])
#
# from autoaugment import ImageNetPolicy
# tra =transforms.Compose([
# transforms.RandomResizedCrop(224),
# transforms.RandomHorizontalFlip(),
# ImageNetPolicy(),
# transforms.ToTensor(),
# normalize])
# image = PIL.Image.open(path)
# policy = ImageNetPolicy()
# transformed = policy(image)
if args.dataset == 'amd':
from datasets.amd_dataset import traindataset
elif args.dataset == 'pm':
from datasets.pm_dataset import traindataset
elif args.dataset == "drdme":
from datasets.drdme_dataset import traindataset
elif args.dataset == "missidor":
from datasets.missidor import traindataset
else:
print("no dataset")
exit(0)
if args.evaluate:
# result = validate(val_loader, model, args)
result = multi_validate(model, test_times, normalize, traindataset,
args)
print("acc_dr, acc_dme, acc_joint", result)
return
val_dataset = traindataset(root=args.data,
mode='val',
transform=tra_test,
num_class=args.num_class,
multitask=args.multitask)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
train_dataset = traindataset(root=args.data,
mode='train',
transform=tra,
num_class=args.num_class,
multitask=args.multitask)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
worker_init_fn=worker_init_fn)
writer = SummaryWriter()
writer.add_text('Text', str(args))
# from lr_scheduler import LRScheduler
# lr_scheduler = LRScheduler(optimizer, len(train_loader), args)
for epoch in range(args.start_epoch, args.epochs):
is_best = False
lr = adjust_learning_rate(optimizer, epoch, args)
writer.add_scalar("lr", lr, epoch)
# train for one epoch
loss_train = train(train_loader, model, criterion, optimizer, args)
writer.add_scalar('Train loss', loss_train, epoch)
# evaluate on validation set
if epoch % 20 == 0:
acc_dr, acc_dme, joint_acc = validate(val_loader, model, args)
writer.add_scalar("Val acc_dr", acc_dr, epoch)
writer.add_scalar("Val acc_dme", acc_dme, epoch)
writer.add_scalar("Val acc_joint", joint_acc, epoch)
is_best = joint_acc >= best_acc1
best_acc1 = max(joint_acc, best_acc1)
if not args.invalid:
if is_best:
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
},
is_best,
filename="checkpoint" + str(epoch) + ".pth.tar",
save_dir=args.model_dir)
def main(args):
flow.enable_eager_execution()
train_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="train",
dataset_size=
7795, # NOTE(Liang Depeng): needs to explictly set the dataset size
batch_size=args.train_batch_size,
)
val_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="val",
dataset_size=7800,
batch_size=args.val_batch_size,
)
# oneflow init
start_t = time.time()
res50_module = resnet50(num_classes=8)
if args.load_checkpoint != "":
res50_module.load_state_dict(flow.load(args.load_checkpoint))
end_t = time.time()
print("init time : {}".format(end_t - start_t))
of_cross_entropy = flow.nn.CrossEntropyLoss()
res50_module.to("cuda")
of_cross_entropy.to("cuda")
of_sgd = flow.optim.SGD(res50_module.parameters(),
lr=args.learning_rate,
momentum=args.mom)
of_losses = []
all_samples = len(val_data_loader) * args.val_batch_size
print_interval = 20
for epoch in range(args.epochs):
res50_module.train()
for b in range(len(train_data_loader)):
image, label = train_data_loader.get_batch()
# oneflow train
start_t = time.time()
image = image.to("cuda")
label = label.to("cuda")
logits, body = res50_module(image)
loss = scloss(label).to("cuda")(body, logits)
loss.backward()
of_sgd.step()
of_sgd.zero_grad()
end_t = time.time()
if b % print_interval == 0:
l = loss.numpy()
of_losses.append(l)
print(
"epoch {} train iter {} oneflow loss {}, train time : {}".
format(epoch, b, l, end_t - start_t))
print("epoch %d train done, start validation" % epoch)
res50_module.eval()
correct_of = 0.0
for b in range(len(val_data_loader)):
image, label = val_data_loader.get_batch()
start_t = time.time()
image = image.to("cuda")
with flow.no_grad():
logits, body = res50_module(image)
predictions = logits.softmax()
of_predictions = predictions.numpy()
clsidxs = np.argmax(of_predictions, axis=1)
label_nd = label.numpy()
for i in range(args.val_batch_size):
if clsidxs[i] == label_nd[i]:
correct_of += 1
end_t = time.time()
print("epoch %d, oneflow top1 val acc: %f" %
(epoch, correct_of / all_samples))
flow.save(
res50_module.state_dict(),
os.path.join(
args.save_checkpoint_path,
"epoch_%d_val_acc_%f" % (epoch, correct_of / all_samples),
),
)
writer = open("of_losses.txt", "w")
for o in of_losses:
writer.write("%f\n" % o)
writer.close()
def setup(args):
train_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="train",
dataset_size=9469,
batch_size=args.train_batch_size,
)
val_data_loader = OFRecordDataLoader(
ofrecord_root=args.ofrecord_path,
mode="val",
dataset_size=3925,
batch_size=args.val_batch_size,
)
criterion = flow.nn.CrossEntropyLoss()
# model setup
eager_model = resnet50()
graph_model = resnet50()
graph_model.load_state_dict(eager_model.state_dict())
eager_model.to("cuda")
graph_model.to("cuda")
# optimizer setup
eager_optimizer = flow.optim.SGD(eager_model.parameters(),
lr=args.learning_rate,
momentum=args.mom)
graph_optimizer = flow.optim.SGD(graph_model.parameters(),
lr=args.learning_rate,
momentum=args.mom)
# criterion setup
criterion = flow.nn.CrossEntropyLoss()
criterion = criterion.to("cuda")
class ModelTrainGraph(flow.nn.Graph):
def __init__(self):
super().__init__()
self.graph_model = graph_model
self.criterion = criterion
self.add_optimizer(graph_optimizer)
def build(self, image, label):
logits = self.graph_model(image)
loss = self.criterion(logits, label)
loss.backward()
return loss
class ModelEvalGraph(flow.nn.Graph):
def __init__(self):
super().__init__()
self.graph_model = graph_model
def build(self, image):
with flow.no_grad():
logits = self.graph_model(image)
predictions = logits.softmax()
return predictions
model_train_graph = ModelTrainGraph()
model_eval_graph = ModelEvalGraph()
dic = {
"train_dataloader": train_data_loader,
"val_dataloader": val_data_loader,
"eager": [eager_model, eager_optimizer, criterion],
"graph": [graph_model, model_train_graph, model_eval_graph],
}
return dic
mtcnn = MTCNN(image_size=160, margin=5, min_face_size=20,
thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
device=device)
detect_faces_and_save(mtcnn,str(folder_name))
#folder name is appended with trimmed,so must take that into account
folder_name_parts = list(folder_name.parts)
print('folder_name_parts',folder_name_parts)
folder_name_parts[-1] = 'trimmed_' + str(folder_name_parts[-1])
folder_name = Path(*folder_name_parts)
#folder_name = Path(folder_name.parent).joinpath('trimmed_'+str(folder_name.name))
folder_name = str(folder_name)
print(folder_name)
#Load the dataset
dataset = LoadDatasetFromFolder(folder_name,parsed_args['class_list_file'][0])
class_mappings = dataset.get_class_mappings()
class_list = dataset.get_classes_list()
X,y = dataset.load()
testdataset = FaceDataset(X,y,transform=True)
print('Test set is of length',testdataset.__len__())
batch_size = 16
#Create the loader
testloader = create_batched_loader(testdataset,batch_size=batch_size)
model = resnet50(num_classes=dataset.num_classes())
#Parse the weight file
file_name = str(Path(parsed_args['weight_file'][0]).absolute())
model = load_init_weights(file_name,model=model,test=True)
model = model.to(device)
loss_fn = nn.CrossEntropyLoss()
#Perform the evaluation
val_evaluation(testloader,model,loss_fn,device,test=True)