def main():
##################################
# Initialize saving directory
##################################
if not os.path.exists(config.save_dir):
os.makedirs(config.save_dir)
##################################
# Logging setting
##################################
logging.basicConfig(
filename=os.path.join(config.save_dir, config.log_name),
filemode='w',
format=
'%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',
level=logging.INFO)
warnings.filterwarnings("ignore")
##################################
# Load dataset
##################################
# train_dataset, validate_dataset = get_trainval_datasets(config.tag, config.image_size)
full_train_dataset = CarDataset('train')
n = len(full_train_dataset)
# train_dataset, validate_dataset = torch.utils.data.random_split(full_train_dataset, [int(n*0.8), n-int(n*0.8)])
train_dataset = full_train_dataset
validate_dataset = full_train_dataset
train_loader, validate_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True,
num_workers=config.workers, pin_memory=True), \
DataLoader(validate_dataset, batch_size=config.batch_size * 4, shuffle=False,
num_workers=config.workers, pin_memory=True)
num_classes = full_train_dataset.num_classes
##################################
# Initialize model
##################################
logs = {}
start_epoch = 0
net = WSDAN(num_classes=num_classes,
M=config.num_attentions,
net=config.net,
pretrained=True)
# feature_center: size of (#classes, #attention_maps * #channel_features)
feature_center = torch.zeros(num_classes, config.num_attentions *
net.num_features).to(device)
if config.ckpt:
# Load ckpt and get state_dict
checkpoint = torch.load(config.ckpt)
# Get epoch and some logs
logs = checkpoint['logs']
start_epoch = int(logs['epoch'])
# Load weights
state_dict = checkpoint['state_dict']
net.load_state_dict(state_dict)
logging.info('Network loaded from {}'.format(config.ckpt))
# load feature center
if 'feature_center' in checkpoint:
feature_center = checkpoint['feature_center'].to(device)
logging.info('feature_center loaded from {}'.format(config.ckpt))
logging.info('Network weights save to {}'.format(config.save_dir))
##################################
# Use cuda
##################################
net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
##################################
# Optimizer, LR Scheduler
##################################
learning_rate = logs['lr'] if 'lr' in logs else config.learning_rate
optimizer = torch.optim.SGD(net.parameters(),
lr=learning_rate,
momentum=0.9,
weight_decay=1e-5)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.9, patience=2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=2,
gamma=0.9)
##################################
# ModelCheckpoint
##################################
callback_monitor = 'val_{}'.format(raw_metric.name)
callback = ModelCheckpoint(savepath=os.path.join(config.save_dir,
config.model_name),
monitor=callback_monitor,
mode='max')
if callback_monitor in logs:
callback.set_best_score(logs[callback_monitor])
else:
callback.reset()
##################################
# TRAINING
##################################
logging.info(
'Start training: Total epochs: {}, Batch size: {}, Training size: {}, Validation size: {}'
.format(config.epochs, config.batch_size, len(train_dataset),
len(validate_dataset)))
logging.info('')
for epoch in range(start_epoch, config.epochs):
callback.on_epoch_begin()
logs['epoch'] = epoch + 1
logs['lr'] = optimizer.param_groups[0]['lr']
logging.info('Epoch {:03d}, Learning Rate {:g}'.format(
epoch + 1, optimizer.param_groups[0]['lr']))
pbar = tqdm(total=len(train_loader), unit=' batches')
pbar.set_description('Epoch {}/{}'.format(epoch + 1, config.epochs))
train(logs=logs,
data_loader=train_loader,
net=net,
feature_center=feature_center,
optimizer=optimizer,
pbar=pbar)
validate(logs=logs, data_loader=validate_loader, net=net, pbar=pbar)
if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
scheduler.step(logs['val_loss'])
else:
scheduler.step()
callback.on_epoch_end(logs, net, feature_center=feature_center)
pbar.close()
def main():
parser = OptionParser()
parser.add_option('-j', '--workers', dest='workers', default=16, type='int',
help='number of data loading workers (default: 16)')
parser.add_option('-e', '--epochs', dest='epochs', default=80, type='int',
help='number of epochs (default: 80)')
parser.add_option('-b', '--batch-size', dest='batch_size', default=16, type='int',
help='batch size (default: 16)')
parser.add_option('-c', '--ckpt', dest='ckpt', default=False,
help='load checkpoint model (default: False)')
parser.add_option('-v', '--verbose', dest='verbose', default=100, type='int',
help='show information for each <verbose> iterations (default: 100)')
parser.add_option('--lr', '--learning-rate', dest='lr', default=1e-3, type='float',
help='learning rate (default: 1e-3)')
parser.add_option('--sf', '--save-freq', dest='save_freq', default=1, type='int',
help='saving frequency of .ckpt models (default: 1)')
parser.add_option('--sd', '--save-dir', dest='save_dir', default='./models/wsdan/',
help='saving directory of .ckpt models (default: ./models/wsdan)')
parser.add_option('--ln', '--log-name', dest='log_name', default='train.log',
help='log name (default: train.log)')
parser.add_option('--mn', '--model-name', dest='model_name', default='model.ckpt',
help='model name (default:model.ckpt)')
parser.add_option('--init', '--initial-training', dest='initial_training', default=1, type='int',
help='train from 1-beginning or 0-resume training (default: 1)')
(options, args) = parser.parse_args()
##################################
# Initialize saving directory
##################################
if not os.path.exists(options.save_dir):
os.makedirs(options.save_dir)
##################################
# Logging setting
##################################
logging.basicConfig(
filename=os.path.join( options.save_dir, options.log_name),
filemode='w',
format='%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',
level=logging.INFO)
warnings.filterwarnings("ignore")
##################################
# Load dataset
##################################
image_size = (256,256)
num_classes = 4
transform = transforms.Compose([transforms.Resize(size=image_size),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])])
train_dataset = CustomDataset(data_root='/mnt/HDD/RFW/train/data/',csv_file='data/RFW_Train40k_Images_Metada.csv',transform=transform)
val_dataset = CustomDataset(data_root='/mnt/HDD/RFW/train/data/',csv_file='data/RFW_Val4k_Images_Metadata.csv',transform=transform)
test_dataset = CustomDataset(data_root='/mnt/HDD/RFW/test/data/',csv_file='data/RFW_Test_Images_Metadata.csv',transform=transform)
train_loader = DataLoader(train_dataset, batch_size=options.batch_size, shuffle=True,num_workers=options.workers, pin_memory=True)
validate_loader = DataLoader(val_dataset, batch_size=options.batch_size * 4, shuffle=False,num_workers=options.workers, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=options.batch_size * 4, shuffle=False,num_workers=options.workers, pin_memory=True)
##################################
# Initialize model
##################################
logs = {}
start_epoch = 0
num_attentions = 32
feature_net = inception_v3(pretrained=True)
net = WSDAN(num_classes=num_classes, M=num_attentions, net='inception_mixed_6e', pretrained=True)
# feature_center: size of (#classes, #attention_maps * #channel_features)
feature_center = torch.zeros(num_classes, num_attentions * net.num_features).to(device)
if options.ckpt:
# Load ckpt and get state_dict
checkpoint = torch.load(options.ckpt)
# Get epoch and some logs
logs = checkpoint['logs']
start_epoch = int(logs['epoch'])
# Load weights
state_dict = checkpoint['state_dict']
net.load_state_dict(state_dict)
logging.info('Network loaded from {}'.format(options.ckpt))
# load feature center
if 'feature_center' in checkpoint:
feature_center = checkpoint['feature_center'].to(device)
logging.info('feature_center loaded from {}'.format(options.ckpt))
logging.info('Network weights save to {}'.format(options.save_dir))
feature_net = inception_v3(pretrained=True)
if options.ckpt:
ckpt = options.ckpt
if options.initial_training == 0:
# Get Name (epoch)
epoch_name = (ckpt.split('/')[-1]).split('.')[0]
start_epoch = int(epoch_name)
# Load ckpt and get state_dict
checkpoint = torch.load(ckpt)
state_dict = checkpoint['state_dict']
# Load weights
net.load_state_dict(state_dict)
logging.info('Network loaded from {}'.format(options.ckpt))
# load feature center
if 'feature_center' in checkpoint:
feature_center = checkpoint['feature_center'].to(torch.device("cuda"))
logging.info('feature_center loaded from {}'.format(options.ckpt))
##################################
# Use cuda
##################################
net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
##################################
# Optimizer, LR Scheduler
##################################
learning_rate = logs['lr'] if 'lr' in logs else options.lr
optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-5)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.9, patience=2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.9)
##################################
# ModelCheckpoint
##################################
callback_monitor = 'val_{}'.format(raw_metric.name)
callback = ModelCheckpoint(savepath=os.path.join(options.save_dir, options.model_name),
monitor=callback_monitor,
mode='max')
if callback_monitor in logs:
callback.set_best_score(logs[callback_monitor])
else:
callback.reset()
##################################
# TRAINING
##################################
logging.info('')
logging.info('Start training: Total epochs: {}, Batch size: {}, Training size: {}, Validation size: {}'.
format(options.epochs, options.batch_size, len(train_dataset), len(val_dataset)))
for epoch in range(start_epoch, options.epochs):
callback.on_epoch_begin()
logs['epoch'] = epoch + 1
logs['lr'] = optimizer.param_groups[0]['lr']
logging.info('Epoch {:03d}, Learning Rate {:g}'.format(epoch + 1, optimizer.param_groups[0]['lr']))
pbar = tqdm(total=len(train_loader), unit=' batches')
pbar.set_description('Epoch {}/{}'.format(epoch + 1, options.epochs))
train(logs=logs,
data_loader=train_loader,
net=net,
feature_center=feature_center,
optimizer=optimizer,
pbar=pbar)
validate(logs=logs,
data_loader=validate_loader,
net=net,
pbar=pbar)
if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
scheduler.step(logs['val_loss'])
else:
scheduler.step()
callback.on_epoch_end(logs, net, feature_center=feature_center)
pbar.close()
def main():
##################################
# Initialize saving directory
##################################
if not os.path.exists(config.save_dir):
os.makedirs(config.save_dir)
##################################
# Logging setting
##################################
logging.basicConfig(
filename=os.path.join(config.save_dir, config.log_name),
filemode='w',
format=
'%(asctime)s: %(levelname)s: [%(filename)s:%(lineno)d]: %(message)s',
level=logging.INFO)
warnings.filterwarnings("ignore")
##################################
# Load dataset
##################################
# train_data_size = int(config.num_image * 0.8)
# val_data_size = config.num_image - train_data_size
indices_train = np.random.RandomState(777).permutation(500)
indices_test = np.random.RandomState(777).permutation(140)
print(indices_test)
print('tran_data_size', len(indices_train), 'val_data_size',
len(indices_test))
train_dataset = CervicalDataset(phase='train',
resize=config.image_size,
indices=indices_train)
validate_dataset = CervicalDataset(phase='val',
resize=config.image_size,
indices=indices_test)
train_loader, validate_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True,
num_workers=config.workers, pin_memory=True), \
DataLoader(validate_dataset, batch_size=config.batch_size * 4, shuffle=False,
num_workers=config.workers, pin_memory=True)
num_classes = train_dataset.num_classes
##################################
# Initialize model
##################################
logs = {}
start_epoch = 0
#net = torch.hub.load('moskomule/senet.pytorch', 'se_resnet50', pretrained=True,)
net = res2net50(pretrained=True, num_classes=2)
print(net)
net.aux_logits = False
# i=0
# for m in net.children():
# if isinstance(m, nn.Conv2d) and i < 3:
# for param in m.parameters():
# param.requires_grad=False
# i = i+1
# print(net)
# feature_center: size of (#classes, #attention_maps * #channel_features)
# feature_center = torch.zeros(num_classes, config.num_attentions * net.num_features).to(device)
if config.ckpt:
# Load ckpt and get state_dict
checkpoint = torch.load(config.ckpt)
# Get epoch and some logs
logs = checkpoint['logs']
start_epoch = int(logs['epoch'])
# Load weights
state_dict = checkpoint['state_dict']
net.load_state_dict(state_dict)
logging.info('Network loaded from {}'.format(config.ckpt))
# load feature center
if 'feature_center' in checkpoint:
feature_center = checkpoint['feature_center'].to(device)
logging.info('feature_center loaded from {}'.format(config.ckpt))
logging.info('Network weights save to {}'.format(config.save_dir))
##################################
# Use cuda
##################################
net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
##################################
# Optimizer, LR Scheduler
##################################
learning_rate = config.learning_rate
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
net.parameters()),
lr=learning_rate,
momentum=0.9,
weight_decay=1e-5)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.9, patience=2)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=2,
gamma=0.9)
#net, optimizer = amp.initialize(net, optimizer, opt_level='O1')
##################################
# ModelCheckpoint
##################################
callback_monitor = 'val_{}'.format(raw_metric.name)
callback = ModelCheckpoint(savepath=os.path.join(config.save_dir,
config.model_name),
monitor=callback_monitor,
mode='max')
if callback_monitor in logs:
callback.set_best_score(logs[callback_monitor])
else:
callback.reset()
##################################
# TRAINING
##################################
logging.info(
'Start training: Total epochs: {}, Batch size: {}, Training size: {}, Validation size: {}'
.format(config.epochs, config.batch_size, len(train_dataset),
len(validate_dataset)))
logging.info('')
for epoch in range(start_epoch, config.epochs):
callback.on_epoch_begin()
logs['epoch'] = epoch + 1
logs['lr'] = optimizer.param_groups[0]['lr']
logging.info('Epoch {:03d}, Learning Rate {:g}'.format(
epoch + 1, optimizer.param_groups[0]['lr']))
pbar = tqdm(total=len(train_loader), unit=' batches')
pbar.set_description('{}/{}'.format(epoch + 1, config.epochs))
train(logs=logs,
data_loader=train_loader,
net=net,
optimizer=optimizer,
pbar=pbar)
validate(logs=logs, data_loader=validate_loader, net=net, pbar=pbar)
if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
scheduler.step(logs['val_loss'])
else:
scheduler.step()
callback.on_epoch_end(logs, net)
torch.save(net.state_dict(),
os.path.join(config.save_dir, config.model_name))
pbar.close()