def training():
train_dataset = TrainDataset(
image_dir=args.train_data_dir,
csv_path_train=
f'data/dataset_csv/list_combined_{args.train_type}_small_train.tsv',
csv_path_val=
f'data/dataset_csv/list_combined_{args.train_type}_small_val.tsv',
train_type=args.train_type,
batch_size=16,
shuffle=True,
random_seed=60,
image_shape=args.input_shape)
fm = FashionModel()
fm.create_model(num_classes=train_dataset.num_classes,
input_shape=[args.input_shape[0], args.input_shape[1], 3])
if args.checkpoint is not None:
fm.model.load_weights(args.checkpoint)
fm.model.summary()
trainer = Trainer(model=fm.model,
train_gen=train_dataset.train_generator,
val_gen=train_dataset.validation_generator,
epoch=args.epoch,
step=args.step)
trainer.train(log_dir=args.log_dir)
def test_net():
data_set=TrainDataset()
classes=data_set.classes
net=SSD(len(classes)+1)
_,_,last_time_model=get_check_point()
# assign directly
# last_time_model='./weights/weights_21_110242'
if os.path.exists(last_time_model):
model=torch.load(last_time_model)
net.load_state_dict(model)
print("Using the model from the last check point:`%s`"%(last_time_model))
else:
raise ValueError("no model existed...")
net.eval()
is_cuda=cfg.use_cuda
did=cfg.device_id
img_src=read_image('./data/img/dog.jpg')
w,h=img_src.size
img=TestTransform(img_src,torch.tensor([[0,0,1,1]]).float()) # [c,h,w]
img=img[None]
if is_cuda:
net.cuda(did)
img=img.cuda(did)
boxes,labels,probs=net.predict(img,torch.tensor([[w,h]]).type_as(img))[0]
prob_mask=probs>cfg.out_thruth_thresh
boxes=boxes[prob_mask ]
labels=labels[prob_mask ].long()
probs=probs[prob_mask]
img_src=np.array(img_src) # [h,w,3] 'RGB'
# change to 'BGR'
img_src=img_src[:,:,::-1].copy()
if len(boxes) !=0:
draw_box(img_src,boxes,color='pred',
text_list=[
classes[_]+'[%.3f]'%(__) for _,__ in zip(labels,probs)
]
)
show_img(img_src,-1)
def test_epoch(config, net, loss, epoch, data_dir):
net.eval()
rotate_factor = args.rotate_limit * epoch / args.epochs
dataset = TrainDataset(data_dir, valid_images, config, rotate_factor)
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
metrics = []
for data, target, coord, image_path in tqdm(data_loader):
data = data.cuda()
target = target.cuda()
coord = coord.cuda()
output = net(data, coord)
loss_output = loss(output, target)
metrics.append(loss_output[:13])
metrics = np.asarray(metrics, np.float32)
loss_total = np.mean(metrics[:, 0])
pos_true = int(np.sum(metrics[:, 6]))
pos_total = int(np.sum(metrics[:, 7]))
neg_true = int(np.sum(metrics[:, 8]))
neg_total = int(np.sum(metrics[:, 9]))
pos_margin = int(np.sum(metrics[:, 11]))
neg_margin = int(np.sum(metrics[:, 12]))
print('-' * 30)
print('Validation Results:')
print(
'Total Loss={:.6f}, True Positive={:d}/{:d}, True Negative={:d}/{:d}'.
format(loss_total, pos_true, pos_total, neg_true, neg_total))
print('Hard Samples: pos={}, neg={}'.format(pos_margin, neg_margin))
return (loss_total, loss_class, loss_regress, tpr, tnr)
def main(args, model):
tr_dataset = TrainDataset(json_dir=args.json_path,
batch_size=args.batch_size)
cv_dataset = CvDataset(json_dir=args.json_path,
batch_size=args.cv_batch_size)
tr_loader = TrainDataLoader(data_set=tr_dataset,
batch_size=1,
num_workers=args.num_workers)
cv_loader = CvDataLoader(data_set=cv_dataset,
batch_size=1,
num_workers=args.num_workers)
data= {'tr_loader': tr_loader, 'cv_loader': cv_loader}
print(model)
# count the parameter number of the network
print('The number of trainable parameters of the net is:%d' % (numParams(model)))
model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.l2)
solver = Solver(data, model, optimizer, args)
solver.train()
A.Cutout(p=0.5),
A.OneOf(
[A.HueSaturationValue(), A.ShiftScaleRotate()], p=1),
A.Normalize(mean=[0.4452, 0.4457, 0.4464], std=[0.2592, 0.2596, 0.2600]),
ToTensorV2(),
])
test_transform = A.Compose([
A.SmallestMaxSize(args.max_size),
#A.CenterCrop(args.image_size, args.image_size, p=1.),
A.Normalize(mean=[0.4452, 0.4457, 0.4464], std=[0.2592, 0.2596, 0.2600]),
ToTensorV2(),
])
# Dataset, Dataloader 정의
dataset = TrainDataset(args, transform=train_transform)
# Add remained last one data
train_size = int(len(dataset) * 0.8) + 1
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
val_dataset.transform = test_transform
test_dataset = TestDataset(args, transform=test_transform)
train_sampler = RandomSampler(train_dataset)
val_sampler = SequentialSampler(val_dataset)
test_sampler = SequentialSampler(test_dataset)
train_loader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.batch_size,
def train():
net=Darknet('./net/yolo/cfg/yolo.cfg','./models/extraction.weights')
net._print()
writer = SummaryWriter('log')
data_set=TrainDataset(net.net_width,net.net_height)
data_loader=DataLoader(
data_set,
batch_size=cfg.batch_size,
shuffle=True,
drop_last=False,
num_workers=cfg.num_worker
)
# test the dataloader...
# for b_imgs,b_boxes,b_labels,b_real_box_num in tqdm(data_loader):
# tqdm.write('%s %s %s'%(str(b_imgs.shape),str(b_imgs.max()),str(b_imgs.min())) )
epoch,iteration,w_path=get_check_point()
if w_path:
model=torch.load(w_path)
net.load_state_dict(model)
print("Using the model from the last check point:%s"%(w_path) )
epoch+=1
net.train()
is_cuda=cfg.use_cuda
did=cfg.device_id
if is_cuda:
net.cuda(did)
print( list(net.children()) )
while epoch<cfg.epochs:
# print('********\t EPOCH %d \t********' % (epoch))
for b_imgs,b_boxes,b_labels,b_real_box_num in tqdm(data_loader):
if is_cuda:
b_imgs=b_imgs.cuda(did)
b_boxes=b_boxes.cuda(did)
b_labels=b_labels.cuda(did)
b_real_box_num=b_real_box_num.cuda(did)
_loss=net(b_imgs,b_boxes,b_labels,b_real_box_num,writer,iteration)
_loss=net.opt_step(_loss)
writer.add_scalar('Train/Loss',_loss,iteration)
tqdm.write('Epoch:%d, iter:%d, loss:%.5f'%(epoch,iteration,_loss))
iteration+=1
adjust_lr(net.optimizer,iteration,cfg.lrs)
if epoch % cfg.save_per_epoch ==0:
torch.save(net.state_dict(),'%sweights_%d_%d'%(cfg.weights_dir,epoch,iteration) )
if epoch % cfg.eval_per_epoch==0:
_map= eval_net(net=net,num=100,shuffle=True)['map']
writer.add_scalar('Train/Eval',_map,iteration)
print("map:",_map)
epoch+=1
print(eval_net(net=net))
writer.close()
parser.add_argument("--gpu", dest="gpu", action="store_true")
parser.add_argument("--cpu", dest="gpu", action="store_false")
parser.set_defaults(gpu=True)
if __name__ == "__main__":
args = parser.parse_args()
device = torch.device("cuda:0" if (
torch.cuda.is_available() and args.gpu) else "cpu")
now = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
results_dir = f"results/{now}"
models_dir = f"models/{now}"
os.makedirs(results_dir, exist_ok=True)
os.makedirs(models_dir, exist_ok=True)
train_dataset = TrainDataset(
root=args.train_data_path,
scale_factor=args.scale_factor,
hr_size=args.hr_size,
random_crop_size=args.random_crop_size,
)
test_dataset = TestDataset(root=args.test_data_path,
scale_factor=args.scale_factor,
hr_size=args.hr_size)
train_dataloader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=64)
gen_net = Generator(num_res_blocks=args.gen_res_blocks,
upscale_factor=args.scale_factor).to(device)
dis_net = Discriminator(hr_size=args.random_crop_size,
sigmoid=not args.no_sigmoid).to(device)
print(f"Generator number of parameters: {count_parameters(gen_net)}")
print(f"Discriminator number of parameters: {count_parameters(dis_net)}")
model_path = os.path.join(os.getcwd(),"models", MODEL_NAME)
model = SRFBN(scale_factor=4, fe_block=FE_BLOCK)
chkpt = torch.load(model_path)
model.load_state_dict(chkpt['state_dict'])
model.to(device)
all_test_dataset_path = os.path.join(os.getcwd(), "test_data")
all_test_dataset = os.listdir(all_test_dataset_path)
# iterating over datasets
for dataset in all_test_dataset:
print(f'Testing on {dataset}')
test_dataset_path = os.path.join(all_test_dataset_path, dataset)
test_dataset = TrainDataset(test_dataset_path)
test_dataloader = data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
model.eval()
avg_psnr = 0.0
avg_ssim = 0.0
max_psnr = 0.0
max_ssim = 0.0
img_name1 = ""
img_name2 = ""
with torch.no_grad():
for idx, batch_data in enumerate(test_dataloader):
lr_image, hr_image = batch_data
lr_image = lr_image.to(device)
lr_image = torch.transpose(lr_image, 3, 1)
from my_utils import *
from light_cnn import LightCNN_29Layers_v2
import torch.nn.functional as F
from torchvision.utils import make_grid
from matplotlib import pyplot as plt
from torchsummary import summary
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
if __name__ == "__main__":
cudnn.enabled = True
img_list = open(config.train['img_list'], 'r').read().split('\n')
img_list.pop()
# input
trainloader = torch.utils.data.DataLoader(TrainDataset(img_list), batch_size=config.train['batch_size'],
shuffle=True, num_workers=2, pin_memory=True)
G = Generator(zdim=config.G['zdim'], use_batchnorm=config.G['use_batchnorm'],
use_residual_block=config.G['use_residual_block'],
num_classes=config.G['num_classes']).cuda()
D = Discriminator(use_batchnorm=True).cuda()
feature_extract_model = LightCNN_29Layers_v2(num_classes=80013).cuda()
feature_extract_model = torch.nn.DataParallel(feature_extract_model).cuda()
feature_extract_model.module.fc2 = torch.nn.Linear(in_features=256, out_features=360).cuda()
torch.nn.init.kaiming_uniform_(feature_extract_model.module.fc2.weight)
one = torch.FloatTensor([1])
mone = one * -1
resume_model(feature_extract_model, 'multipie_finetune_model')
feature_extract_model.eval()
def train_epoch(config, net, loss, optimizer, epoch, data_dir):
net.train()
rotate_factor = args.rotate_limit * epoch / args.epochs
dataset = TrainDataset(data_dir, train_images, config, rotate_factor)
data_loader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
lr = get_lr(epoch)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
metrics = []
result_epoch = []
epoch_fmt = 'Epoch {}/{} training for {} exmaples with lr={:.3e} and rotate_factor={:.1f}'
print(
epoch_fmt.format(epoch, args.epochs, len(data_loader), lr,
rotate_factor))
for data, target, coord, image_path in tqdm(data_loader):
# print(np.where(target[...,0]>0.5))
pos_in_image = np.where(target[..., 0] > 0.5)[0]
data = data.cuda()
target = target.cuda()
coord = coord.cuda()
output = net(data, coord)
loss_output = loss(output, target)
optimizer.zero_grad()
loss_output[0].backward()
optimizer.step()
loss_output[0] = loss_output[0].item()
metrics.append(loss_output[:10])
coord_start = coord[:, :, 0, 0, 0].detach().cpu().numpy()
coord_end = coord[:, :, -1, -1, -1].detach().cpu().numpy()
# print('coord: start =', coord_start, ', end =', coord_end)
result_batch = {
'epoch': epoch,
}
result_batch['loss'] = loss_output[0]
result_batch['loss_ce'] = loss_output[1]
result_batch['delta_z'] = loss_output[2]
result_batch['delta_h'] = loss_output[3]
result_batch['delta_w'] = loss_output[4]
result_batch['delta_d'] = loss_output[5]
result_batch['pos_true'] = 0 if loss_output[7] == 0 else int(
loss_output[6])
result_batch['pos_false'] = loss_output[7] - result_batch['pos_true']
result_batch['pos_total'] = loss_output[7]
result_batch['neg_true'] = 0 if loss_output[9] == 0 else int(
loss_output[8])
result_batch['neg_false'] = loss_output[9] - result_batch['neg_true']
result_batch['neg_total'] = loss_output[9]
result_batch['num_total'] = loss_output[10]
result_batch['pos_margin'] = loss_output[11]
result_batch['neg_margin'] = loss_output[12]
for i_image in range(len(image_path)):
result_batch['image_{}'.format(i_image)] = image_path[i_image]
result_batch['coord_{}'.format(i_image)] = (
list(coord_start)[i_image], list(coord_end)[i_image])
for i, i_pos in enumerate(pos_in_image):
result_batch['pred_{}'.format(i_pos)] = loss_output[13 + i]
result_epoch.append(result_batch)
if epoch % args.save_freq == 0:
state_dict = net.module.state_dict()
for key in state_dict.keys():
state_dict[key] = state_dict[key].cpu()
torch.save(
{
'epoch': epoch,
'save_dir': args.output,
'state_dict': state_dict,
'args': args
}, os.path.join(args.output, '%03d.ckpt' % epoch))
metrics = np.asarray(metrics, np.float32)
loss_total = np.mean(metrics[:, 0])
loss_class = np.mean(metrics[:, 1])
loss_regress = [
np.mean(metrics[:, 2]),
np.mean(metrics[:, 3]),
np.mean(metrics[:, 4]),
np.mean(metrics[:, 5])
]
tpr = 100.0 * np.sum(metrics[:, 6]) / np.sum(metrics[:, 7])
tnr = 100.0 * np.sum(metrics[:, 8]) / np.sum(metrics[:, 9])
df_results = pd.DataFrame(result_epoch)
df_class = df_results[[
'pos_true', 'pos_false', 'pos_total', 'neg_false', 'neg_true',
'neg_total', 'num_total'
]]
df_regre = df_results[[
'loss_ce', 'delta_z', 'delta_h', 'delta_w', 'delta_d'
]]
print('Training results:')
print('-' * 30)
print(df_class.sum())
print('-' * 20)
print(df_regre.mean())
return (loss_total, loss_class, loss_regress, tpr, tnr), result_epoch
feature_extract_model = torch.nn.DataParallel(feature_extract_model).cuda()
feature_extract_model.module.fc2 = nn.Linear(in_features=256,
out_features=360).cuda()
torch.nn.init.kaiming_uniform_(feature_extract_model.module.fc2.weight)
optim_LCNN = torch.optim.Adam(
feature_extract_model.parameters(),
lr=1e-4,
)
resume_model(feature_extract_model, 'model_save')
resume_optimizer(optim_LCNN, feature_extract_model, 'model_save')
# Train LightCNN on multipie
# input
trainloader = torch.utils.data.DataLoader(TrainDataset(img_list),
batch_size=64,
shuffle=True,
num_workers=2,
pin_memory=True)
last_epoch = 2
cross_entropy = nn.CrossEntropyLoss().cuda()
epoch = 80
for epoch in range(last_epoch + 1, 10):
top1 = AverageMeter()
for step, batch in enumerate(trainloader):
for k in batch:
if k != 'name':
batch[k] = Variable(batch[k].cuda(async=True),
import torch
from torch import nn
from torch import optim
from torch.utils import data
from data import TrainDataset, TestDataset
from model import Model
from logger import Logger
import os
batch_size = 16
end_epoch = 500
load_checkpoint = False
save_path = '/media/disk1/pxq/ASD/' # path to save checkpoint
train_dataset = TrainDataset()
train_loader = data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = TestDataset()
test_loader = data.DataLoader(test_dataset, batch_size=1, shuffle=False)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model().to(device)
criterion = nn.MSELoss(size_average=False).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
# log to tensorboard
loss_logger = Logger('./logs/loss')
eval_logger = Logger('./logs/eval')
# load checkpoint
def __getitem__(self, idx):
img, box, label, scale, cur_img_size = TrainDataset.__getitem__(
self, idx)
return img, box, label, scale, idx, cur_img_size
def train():
cfg._print()
# im=torch.randn([1,3,300,300])
# for loc,cls in net(im):
# print(loc.shape,cls.shape)
# for k,v in net.named_parameters():
# print(k,v,v.shape)
# boxes=get_default_boxes()
# print(boxes.shape)
data_set=TrainDataset()
data_loader=DataLoader(
data_set,
batch_size=cfg.batch_size,
shuffle=True,
drop_last=False,
num_workers=cfg.num_worker
)
# NOTE: plus one, en?
net=SSD(len(data_set.classes)+1)
net._print()
epoch,iteration,w_path=get_check_point()
if w_path:
model=torch.load(w_path)
net.load_state_dict(model)
print("Using the model from the last check point:%s"%(w_path) )
epoch+=1
net.train()
is_cuda=cfg.use_cuda
did=cfg.device_id
if is_cuda:
net.cuda(did)
while epoch<cfg.epochs:
# print('********\t EPOCH %d \t********' % (epoch))
for i,(imgs,targets,labels) in tqdm(enumerate(data_loader)):
if is_cuda:
imgs=imgs.cuda(did)
targets=targets.cuda(did)
labels=labels.cuda(did)
_loss=net.train_once(imgs,targets,labels)
tqdm.write('Epoch:%d, iter:%d, loss:%.5f'%(epoch,iteration,_loss))
iteration+=1
adjust_lr(net.optimizer,iteration,cfg.lrs)
if epoch % cfg.save_per_epoch ==0:
torch.save(net.state_dict(),'%sweights_%d_%d'%(cfg.weights_dir,epoch,iteration) )
if epoch % cfg.eval_per_epoch==0:
_map= eval_net(net=net,num=100,shuffle=True)['map']
print("map:",_map)
epoch+=1
print(eval_net(net=net))
def main(args):
if (not args.do_train) and (not args.do_valid) and (not args.do_test):
raise ValueError('one of train/val/test mode must be choosed.')
if args.init_checkpoint:
override_config(args)
elif args.data_path is None:
raise ValueError('one of init_checkpoint/data_path must be choosed.')
if args.do_train and args.save_path is None:
raise ValueError('Where do you want to save your trained model?')
if args.save_path and not os.path.exists(args.save_path):
os.makedirs(args.save_path)
# Write logs to checkpoint and console
set_logger(args)
data_reader = DataReader(args.data_path)
num_entity = len(data_reader.entity_dict)
num_relation = len(data_reader.relation_dict)
logging.info('Model: {}'.format(args.model))
logging.info('Data Path: {}'.format(args.data_path))
logging.info('Num Entity: {}'.format(num_entity))
logging.info('Num Relation: {}'.format(num_relation))
logging.info('Num Train: {}'.format(len(data_reader.train_data)))
logging.info('Num Valid: {}'.format(len(data_reader.valid_data)))
logging.info('Num Test: {}'.format(len(data_reader.test_data)))
if args.model == 'ModE':
kge_model = ModE(num_entity, num_relation, args.hidden_dim, args.gamma)
elif args.model == 'HAKE':
kge_model = HAKE(num_entity, num_relation, args.hidden_dim, args.gamma, args.modulus_weight, args.phase_weight)
logging.info('Model Parameter Configuration:')
for name, param in kge_model.named_parameters():
logging.info('Parameter %s: %s, require_grad = %s' % (name, str(param.size()), str(param.requires_grad)))
kge_model = kge_model.cuda()
if args.do_train:
# Set training dataloader iterator
train_dataloader_head = DataLoader(
TrainDataset(data_reader, args.negative_sample_size, BatchType.HEAD_BATCH),
batch_size=args.batch_size,
shuffle=True,
num_workers=max(1, args.cpu_num // 2),
collate_fn=TrainDataset.collate_fn
)
train_dataloader_tail = DataLoader(
TrainDataset(data_reader, args.negative_sample_size, BatchType.TAIL_BATCH),
batch_size=args.batch_size,
shuffle=True,
num_workers=max(1, args.cpu_num // 2),
collate_fn=TrainDataset.collate_fn
)
train_iterator = BidirectionalOneShotIterator(train_dataloader_head, train_dataloader_tail)
# Set training configuration
current_learning_rate = args.learning_rate
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, kge_model.parameters()),
lr=current_learning_rate
)
warm_up_steps = args.max_steps // 2
if args.init_checkpoint:
# Restore model from checkpoint directory
logging.info('Loading checkpoint %s...' % args.init_checkpoint)
checkpoint = torch.load(os.path.join(args.init_checkpoint, 'checkpoint'))
init_step = checkpoint['step']
kge_model.load_state_dict(checkpoint['model_state_dict'])
if args.do_train:
current_learning_rate = checkpoint['current_learning_rate']
warm_up_steps = checkpoint['warm_up_steps']
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
else:
logging.info('Randomly Initializing %s Model...' % args.model)
init_step = 1
step = init_step
logging.info('Start Training...')
logging.info('init_step = %d' % init_step)
if not args.do_test:
logging.info('learning_rate = %d' % current_learning_rate)
logging.info('batch_size = %d' % args.batch_size)
logging.info('hidden_dim = %d' % args.hidden_dim)
logging.info('gamma = %f' % args.gamma)
logging.info('adversarial_temperature = %f' % args.adversarial_temperature)
if args.do_train:
training_logs = []
# Training Loop
for step in range(init_step, args.max_steps):
log = kge_model.train_step(kge_model, optimizer, train_iterator, args)
training_logs.append(log)
if step >= warm_up_steps:
if not args.no_decay:
current_learning_rate = current_learning_rate / 10
logging.info('Change learning_rate to %f at step %d' % (current_learning_rate, step))
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, kge_model.parameters()),
lr=current_learning_rate
)
warm_up_steps = warm_up_steps * 3
if step % args.save_checkpoint_steps == 0:
save_variable_list = {
'step': step,
'current_learning_rate': current_learning_rate,
'warm_up_steps': warm_up_steps
}
save_model(kge_model, optimizer, save_variable_list, args)
if step % args.log_steps == 0:
metrics = {}
for metric in training_logs[0].keys():
metrics[metric] = sum([log[metric] for log in training_logs]) / len(training_logs)
log_metrics('Training average', step, metrics)
training_logs = []
if args.do_valid and step % args.valid_steps == 0:
logging.info('Evaluating on Test Dataset...')
metrics = kge_model.test_step(kge_model, data_reader, ModeType.TEST, args)
log_metrics('Valid', step, metrics)
save_variable_list = {
'step': step,
'current_learning_rate': current_learning_rate,
'warm_up_steps': warm_up_steps
}
save_model(kge_model, optimizer, save_variable_list, args)
if args.do_valid:
logging.info('Evaluating on Valid Dataset...')
metrics = kge_model.test_step(kge_model, data_reader, ModeType.VALID, args)
log_metrics('Valid', step, metrics)
if args.do_test:
logging.info('Evaluating on Test Dataset...')
metrics = kge_model.test_step(kge_model, data_reader, ModeType.TEST, args)
log_metrics('Test', step, metrics)
import time
from log import TensorBoardX
from utils import *
import feature_extract_network
import importlib
test_time = False
if __name__ == "__main__":
# img_list = open(config.train['img_list'],'r').read().split('\n')
# img_list.pop()
#input
img_dir = config.train['img_dir']
dataloader = torch.utils.data.DataLoader(
TrainDataset(img_dir),
batch_size=config.train['batch_size'],
shuffle=True,
num_workers=8,
pin_memory=True)
# 为了使用多个gpu
G = torch.nn.DataParallel(
Generator(zdim=config.G['zdim'],
use_batchnorm=config.G['use_batchnorm'],
use_residual_block=config.G['use_residual_block'],
num_classes=config.G['num_classes'])).cuda()
D = torch.nn.DataParallel(
Discriminator(use_batchnorm=config.D['use_batchnorm'])).cuda()
optimizer_G = torch.optim.Adam(filter(lambda p: p.requires_grad,
G.parameters()),
lr=1e-4)
model.load_state_dict(chkpt['state_dict'])
model = nn.DataParallel(model, device_ids=[GPU_ID])
optimizer.load_state_dict(chkpt['optimizer'])
else:
model = nn.DataParallel(model, device_ids=[GPU_ID])
model.apply(init_weights)
# L1 Loss
criterion = nn.L1Loss()
# Learning Rate Scheduler, halves learning rate after 200 epochs
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.5)
model.to(device)
# get training data -> create a dataloader so that complete data is not loaded into memory
train_dataset = TrainDataset(train_data_path, mode=MODE)
train_dataloader = data.DataLoader(train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=NUM_WORKERS)
write_data = []
val_dataset = TrainDataset(val_data_path)
val_dataloader = data.DataLoader(val_dataset,
batch_size=VAL_BATCH_SIZE,
shuffle=False)
last_psnr = 0
last_ssim = 0
def train():
print("my name is van")
# let the random counld be the same
if cfg.train_use_offline_feat:
data_set = TrainSetExt()
else:
data_set = TrainDataset()
data_loader = DataLoader(data_set,
batch_size=1,
shuffle=True,
drop_last=False)
net = MyNet(data_set.classes)
net.print_net_info()
epoch, iteration, w_path = get_check_point()
if w_path:
model = torch.load(w_path)
if cfg.train_use_offline_feat:
net.load_state_dict(model)
else:
net.load_state_dict(model)
print("Using the model from the last check point:%s" % (w_path),
end=" ")
epoch += 1
net.train()
is_cuda = cfg.use_cuda
did = cfg.device_id
if is_cuda:
net.cuda(did)
while epoch < cfg.epochs:
# train the rpn
print('******epoch %d*********' % (epoch))
for i, (imgs, boxes, labels, scale,
img_sizes) in tqdm(enumerate(data_loader)):
if is_cuda:
imgs = imgs.cuda(did)
labels = labels.cuda(did)
boxes = boxes.cuda(did)
scale = scale.cuda(did).float()
img_sizes = img_sizes.cuda(did).float()
loss = net.train_once(imgs, boxes, labels, scale, img_sizes,
cfg.train_use_offline_feat)
tqdm.write('Epoch:%d, iter:%d, loss:%.5f' %
(epoch, iteration, loss))
iteration += 1
if cfg.train_use_offline_feat:
torch.save(net.state_dict(),
'%sweights_%d_%d' % (cfg.weights_dir, epoch, iteration))
else:
torch.save(net.state_dict(),
'%sweights_%d_%d' % (cfg.weights_dir, epoch, iteration))
_map = eval_net(net=net, num=100, shuffle=True)['map']
print("map:", _map)
epoch += 1
print(eval_net(net=net))
from torch.autograd import Variable
import time
from log import TensorBoardX
from utils import *
import feature_extract_network
import importlib
test_time = False
if __name__ == "__main__":
img_list = open(config.train['img_list'], 'r').read().split('\n')
img_list.pop()
#input
dataloader = torch.utils.data.DataLoader(
TrainDataset(img_list),
batch_size=config.train['batch_size'],
shuffle=True,
num_workers=8,
pin_memory=True)
G = torch.nn.DataParallel(
Generator(zdim=config.G['zdim'],
use_batchnorm=config.G['use_batchnorm'],
use_residual_block=config.G['use_residual_block'],
num_classes=config.G['num_classes'])).cuda()
D = torch.nn.DataParallel(
Discriminator(use_batchnorm=config.D['use_batchnorm'])).cuda()
optimizer_G = torch.optim.Adam(filter(lambda p: p.requires_grad,
G.parameters()),
lr=1e-4)
lr = 0.1 * lr_rate
else:
lr = 0.01 * lr_rate
return lr
logfile = "./train_valid.log"
sys.stdout = Logger(logfile)
# resume_path = "./checkpoint/20_epoch.pth"
# if resume_path:
# if os.path.isfile(resume_path):
# print("=> loading checkpoint '{}'".format(resume_path))
# checkpoint = torch.load(resume_path)
# model.load_state_dict(checkpoint)
# sets.phase = 'train'
train_data_train = TrainDataset()
train_data_loader = DataLoader(dataset=train_data_train,
batch_size=4,
shuffle=True,
num_workers=12)
test_data_test = TestDataset()
test_data_loader = DataLoader(dataset=test_data_test,
batch_size=2,
shuffle=False,
num_workers=12)
num_epochs = 101
log = logging.getLogger()
train(train_data_loader, test_data_loader, model, optimizer, get_lr, log)