def train_detector(model,
dataset,
cfg,
validate=False,
timestamp=None,
meta=None):
logger = Logging.getLogger()
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [build_dataloader(ds, data=cfg.data) for ds in dataset]
if torch.cuda.is_available():
model = model.cuda(cfg.gpu_ids[0])
model.device = cfg.gpu_ids[0]
if torch.cuda.device_count() > 1:
model = DataParallel(model, device_ids=cfg.gpu_ids)
else:
model.device = 'cpu'
# build runner
optimizer = cfg.optimizer
if 'ema' in cfg:
ema = cfg.ema
else:
ema = None
runner = Runner(model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta,
ema=ema)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# register eval hooks 需要放在日志前面，不然打印不出日志。
if validate:
cfg.data.val.train = False
val_dataset = build_from_dict(cfg.data.val, DATASET)
val_dataloader = build_dataloader(val_dataset,
shuffle=False,
data=cfg.data)
eval_cfg = cfg.get('evaluation', {})
from yolodet.models.hooks.eval_hook import EvalHook
runner.register_hook(EvalHook(val_dataloader, **eval_cfg))
# register hooks
# runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,cfg.checkpoint_config)
runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
cfg.checkpoint_config, cfg.log_config)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def __init__(self):
self.net = ET_Net()
if (ARGS['gpu']):
self.net = DataParallel(module=self.net.cuda())
self.net.load_state_dict(torch.load(ARGS['weight']))
self.metric_dataset = get_dataset(dataset_name=ARGS['dataset'],
part='metric')
def __init__(self, input_size, n_channels, hparams):
self.hparams = hparams
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# define the models
self.model = WaveNet(n_channels=n_channels).to(self.device)
summary(self.model, (input_size, n_channels))
# self.model.half()
if torch.cuda.device_count() > 1:
print("Number of GPUs will be used: ",
torch.cuda.device_count() - 3)
self.model = DP(self.model,
device_ids=list(
range(torch.cuda.device_count() - 3)))
else:
print('Only one GPU is available')
self.metric = Metric()
self.num_workers = 1
########################## compile the model ###############################
# define optimizer
self.optimizer = torch.optim.Adam(params=self.model.parameters(),
lr=self.hparams['lr'],
weight_decay=1e-5)
# weights = torch.Tensor([0.025,0.033,0.039,0.046,0.069,0.107,0.189,0.134,0.145,0.262,1]).cuda()
self.loss = nn.BCELoss() # CompLoss(self.device)
# define early stopping
self.early_stopping = EarlyStopping(
checkpoint_path=self.hparams['checkpoint_path'] + '/checkpoint.pt',
patience=self.hparams['patience'],
delta=self.hparams['min_delta'],
)
# lr cheduler
self.scheduler = ReduceLROnPlateau(
optimizer=self.optimizer,
mode='max',
factor=0.2,
patience=3,
verbose=True,
threshold=self.hparams['min_delta'],
threshold_mode='abs',
cooldown=0,
eps=0,
)
self.seed_everything(42)
self.threshold = 0.75
self.scaler = torch.cuda.amp.GradScaler()
def __init__(self, base_kernel, device_ids, output_device=None, **kwargs):
DataParallel.__init__(self,
module=base_kernel,
device_ids=device_ids,
output_device=output_device,
dim=-2)
self.output_device = output_device if output_device else device_ids[0]
self.__cached_x1 = torch.empty(1)
self.__cached_x2 = torch.empty(1)
def __init__(self, opt, stage0=False, stage1=False, stage2=False):
super(ModelStage, self).__init__(opt)
# ------------------------------------
# define network
# ------------------------------------
self.stage0 = stage0
self.stage1 = stage1
self.stage2 = stage2
self.netG = define_G(opt, self.stage0, self.stage1,
self.stage2).to(self.device)
self.netG = DataParallel(self.netG)
def __init__(self, opt):
super(ModelGAN, self).__init__(opt)
# ------------------------------------
# define network
# ------------------------------------
self.netG = define_G(opt).to(self.device)
self.netG = DataParallel(self.netG)
if self.is_train:
self.netF = define_F(opt).to(self.device)
self.netD = define_D(opt).to(self.device)
self.netF = DataParallel(self.netF)
self.netD = DataParallel(self.netD)
def __init__(self, opt):
super(Ranker_Model, self).__init__(opt)
if opt['dist']:
self.rank = torch.distributed.get_rank()
else:
self.rank = -1 # non dist training
train_opt = opt['train']
# define networks and load pretrained models
self.netR = networks.define_R(opt).to(self.device)
if opt['dist']:
self.netR = DistributedDataParallel(self.netR, device_ids=[torch.cuda.current_device()])
else:
self.netR = DataParallel(self.netR)
self.load()
if self.is_train:
self.netR.train()
# loss
self.RankLoss = nn.MarginRankingLoss(margin=0.5)
self.RankLoss.to(self.device)
self.L2Loss = nn.L1Loss()
self.L2Loss.to(self.device)
# optimizers
self.optimizers = []
wd_R = train_opt['weight_decay_R'] if train_opt['weight_decay_R'] else 0
optim_params = []
for k, v in self.netR.named_parameters(): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
else:
print('WARNING: params [%s] will not optimize.' % k)
self.optimizer_R = torch.optim.Adam(optim_params, lr=train_opt['lr_R'], weight_decay=wd_R)
print('Weight_decay:%f' % wd_R)
self.optimizers.append(self.optimizer_R)
# schedulers
self.schedulers = []
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
train_opt['lr_steps'], train_opt['lr_gamma']))
else:
raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
print('---------- Model initialized ------------------')
self.print_network()
print('-----------------------------------------------')
def __init__(self, model_weight_file, sys_device_ids=''):
if len(sys_device_ids) > 0:
os.environ['CUDA_VISIBLE_DEVICES'] = sys_device_ids
self.sys_device_ids = sys_device_ids
self.model = DataParallel(Model())
if torch.cuda.is_available() and self.sys_device_ids != '':
device = torch.device('cuda')
else:
device = torch.device('cpu')
self.model.load_state_dict(torch.load(model_weight_file,
map_location=device))
self.model.to(device)
self.model.eval()
def set_device(self, gpus, chunk_sizes, device):
print("Set device: ", gpus, device)
if len(gpus) > 1:
self.model_with_loss = DataParallel(self.model_with_loss,
device_ids=gpus).to(device)
else:
# print("Push model to cpu or one gpu ")
self.model_with_loss = self.model_with_loss.to(device)
for state in self.optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(device=device, non_blocking=True)
class TestProcess:
def __init__(self):
self.net = ET_Net()
if (ARGS['gpu']):
self.net = DataParallel(module=self.net.cuda())
self.net.load_state_dict(torch.load(ARGS['weight']))
self.test_dataset = get_dataset(dataset_name=ARGS['dataset'], part='test')
def predict(self):
start = time.time()
self.net.eval()
test_dataloader = DataLoader(self.test_dataset, batch_size=1) # only support batch size = 1
os.makedirs(ARGS['prediction_save_folder'], exist_ok=True)
for items in test_dataloader:
images, mask, filename = items['image'], items['mask'], items['filename']
images = images.float()
mask = mask.long()
print('image shape:', images.size())
image_patches, big_h, big_w = get_test_patches(images, ARGS['crop_size'], ARGS['stride_size'])
test_patch_dataloader = DataLoader(image_patches, batch_size=ARGS['batch_size'], shuffle=False, drop_last=False)
test_results = []
print('Number of batches for testing:', len(test_patch_dataloader))
for patches in test_patch_dataloader:
if ARGS['gpu']:
patches = patches.cuda()
with torch.no_grad():
result_patches_edge, result_patches = self.net(patches)
test_results.append(result_patches.cpu())
test_results = torch.cat(test_results, dim=0)
# merge
test_results = recompone_overlap(test_results, ARGS['crop_size'], ARGS['stride_size'], big_h, big_w)
test_results = test_results[:, 1, :images.size(2), :images.size(3)] * mask
test_results = Image.fromarray(test_results[0].numpy())
test_results.save(os.path.join(ARGS['prediction_save_folder'], filename[0]))
print(f'Finish prediction for {filename[0]}')
finish = time.time()
print('Predicting time consumed: {:.2f}s'.format(finish - start))
def modelDeploy(args, model, optimizer, scheduler, logger):
if args.num_gpus >= 1:
from torch.nn.parallel import DataParallel
model = DataParallel(model)
model = model.cuda()
if torch.backends.cudnn.is_available():
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
cudnn.deterministic = True
trainData = {'epoch': 0, 'loss': [], 'miou': [], 'val': [], 'bestMiou': 0}
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume,
map_location=torch.device('cpu'))
# model&optimizer
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
# stop point
trainData = checkpoint['trainData']
for i in range(trainData['epoch']):
scheduler.step()
# print(trainData)
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, trainData['epoch']))
else:
logger.error("=> no checkpoint found at '{}'".format(args.resume))
assert False, "=> no checkpoint found at '{}'".format(args.resume)
if args.finetune:
if os.path.isfile(args.finetune):
logger.info("=> finetuning checkpoint '{}'".format(args.finetune))
state_all = torch.load(args.finetune, map_location='cpu')['model']
state_clip = {} # only use backbone parameters
# print(model.state_dict().keys())
for k, v in state_all.items():
state_clip[k] = v
# print(state_clip.keys())
model.load_state_dict(state_clip, strict=False)
else:
logger.warning("finetune is not a file.")
pass
if args.freeze_bn:
logger.warning('Freezing batch normalization layers')
for m in model.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
m.weight.requires_grad = False
m.bias.requires_grad = False
return model, trainData
class VGGLoss(nn.Module):
def __init__(self):
super().__init__()
self.layer_ids = (2, 7, 12, 21, 30)
self.vgg = DataParallel(VGGFeatureExtractor(layer_ids=self.layer_ids))
self.vgg.eval()
self.l1_loss = nn.L1Loss(reduction='mean')
def forward(self, input, target):
features_input = self.vgg(input)
features_target = self.vgg(target)
perceptual_loss = 0
for i in range(len(self.layer_ids)):
perceptual_loss += self.l1_loss(features_input[i], features_target[i]) / len(self.layer_ids)
return perceptual_loss
def __init__(self,
batch=8,
subdivisions=4,
epochs=100,
burn_in=1000,
steps=[400000, 450000]):
_model = build_from_dict(model, DETECTORS)
self.model = DataParallel(_model.cuda(), device_ids=[0])
self.train_dataset = build_from_dict(data_cfg['train'], DATASET)
self.val_dataset = build_from_dict(data_cfg['val'], DATASET)
self.burn_in = burn_in
self.steps = steps
self.epochs = epochs
self.batch = batch
self.subdivisions = subdivisions
self.train_size = len(self.train_dataset)
self.val_size = len(self.val_dataset)
self.train_loader = DataLoader(self.train_dataset,
batch_size=batch // subdivisions,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True,
collate_fn=self.collate)
self.val_loader = DataLoader(self.val_dataset,
batch_size=batch // subdivisions,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True,
collate_fn=self.collate)
self.optimizer = optim.Adam(
self.model.parameters(),
lr=0.001 / batch,
betas=(0.9, 0.999),
eps=1e-08,
)
self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer,
self.burnin_schedule)
def test_data_parallel():
shutil.rmtree(out_dir, ignore_errors=True)
hook = smd.Hook(
out_dir=out_dir,
save_config=smd.SaveConfig(save_steps=[0, 1, 5]),
save_all=True,
include_workers="one",
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Net().to(device)
if device == "cuda":
model = DataParallel(model)
hook.register_module(model)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
train(model, hook, torch.device(device), optimizer, num_steps=10)
trial = create_trial(out_dir)
assert trial.steps() == [0, 1, 5]
if device == "cpu":
assert len(trial.tensor_names()) == 38
else:
assert len(trial.tensor_names()) > 37
shutil.rmtree(out_dir, ignore_errors=True)
def parallelize_submodules(parent_module):
"""
Wraps a parent's sub-modules with a DataParallel object that permits multi
GPU usage.
"""
for mod_name, module in parent_module.named_children():
setattr(parent_module, mod_name, DataParallel(module))
def main():
TVT, TMO = set_devices((-1,))
model = Model(local_conv_out_channels=128,
num_classes=751)
model_w = DataParallel(model)
map_location = (lambda storage, loc: storage)
sd = torch.load('ckpt.pth', map_location=map_location)
load_state_dict(model, sd)
print('Loaded model weights')
ims = []
im1 = PIL.Image.open('demo/0.jpg')
ims.append(preprocess_image(im1))
im2 = PIL.Image.open('demo/1.jpg')
ims.append(preprocess_image(im2))
ims = Variable(TVT(torch.from_numpy(np.array(ims)).float()))
start = time.time()
global_feat, local_feat = model(ims)[:2]
global_feat = np.vstack(global_feat.data.cpu().numpy())
local_feat = local_feat.data.cpu().numpy()
distance = get_distance((global_feat[0],local_feat[0]),(global_feat[1],local_feat[1]))
print("distance: ", distance)
def start():
###############
# preparation #
###############
model = Model(last_conv_stride=last_conv_stride,
num_stripes=num_stripes,
local_conv_out_channels=local_conv_out_channels)
model_w = DataParallel(model)
TMO([model])
# preprocessing
preprocessor = PreProcessIm(resize_h_w=resize_h_w,
scale=scale_im,
im_mean=im_mean,
im_std=im_std)
# load model
map_location = (lambda storage, loc: storage)
sd = torch.load(model_weight_file, map_location=map_location)
load_state_dict(model, sd['state_dicts'][0])
print('Loaded model weight from {}'.format(model_weight_file))
extractor = ExtractFeature(model_w, TVT)
return preprocessor, extractor
def parallelize(self):
if self.args.distributed:
self.model.G = DistributedDataParallel(
self.model.G,
device_ids=[self.args.rank],
output_device=self.args.rank)
if self.model.D is not None:
self.model.D = DistributedDataParallel(
self.model.D,
device_ids=[self.args.rank],
output_device=self.args.rank)
else:
self.model.G = DataParallel(self.model.G, range(self.n_GPUs))
if self.model.D is not None:
self.model.D = DataParallel(self.model.D, range(self.n_GPUs))
def test_is_parallel_module():
class Model(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(2, 2, 1)
def forward(self, x):
return self.conv(x)
model = Model()
assert not is_parallel_module(model)
dp = DataParallel(model)
assert is_parallel_module(dp)
mmdp = MMDataParallel(model)
assert is_parallel_module(mmdp)
ddp = DistributedDataParallel(model, process_group=MagicMock())
assert is_parallel_module(ddp)
mmddp = MMDistributedDataParallel(model, process_group=MagicMock())
assert is_parallel_module(mmddp)
deprecated_mmddp = DeprecatedMMDDP(model)
assert is_parallel_module(deprecated_mmddp)
def __init__(self, config):
self.arch = config['model']['arch']
self.heads = config['model']['heads']
self.head_conv = config['model']['head_conv']
self.lr = config['train']['lr']
self.lr_step = config['train']['lr_step']
self.resume_path = config['train']['resume_path']
self.seed = config['train']['seed']
self.not_cuda_benchmark = config['train']['not_cuda_benchmark']
self.gpu_str = config['gpu_str']
self.gpus = config['gpus']
self.chunk_sizes = config['train']['chunk_sizes']
self.num_epochs = config['train']['num_epochs']
self.batch_size = config['train']['batch_size']
self.num_workers = config['train']['num_workers']
self.val_intervals = config['train']['val_intervals']
self.save_all = config['train']['save_all']
self.metric = config['train']['metric']
self.save_dir = config['save_dir']
self.num_iters = config['train']['num_iters']
self.exp_id = config['exp_id']
self.print_iter = config['train']['print_iter']
self.hide_data_time = config['train']['hide_data_time']
self.model = create_model(self.arch, self.heads, self.head_conv)
self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr)
self.loss_stats, self.loss = self._get_losses(config)
self.model_with_loss = ModelWithLoss(self.model, self.loss)
if len(self.gpus) == 1 and self.gpu_str != '-1':
torch.cuda.set_device(int(self.gpus[0]))
self.device = torch.device('cuda' if torch.cuda.is_available()
and self.gpu_str != '-1' else 'cpu')
self.set_device(self.gpus, self.chunk_sizes, self.device)
self.start_epoch = 0
if self.resume_path != '':
self.model, self.optimizer, self.start_epoch = load_model(
self.model, self.resume_path, self.optimizer, self.lr,
self.lr_step)
self.train_loader = torch.utils.data.DataLoader(
DATASET_CUSTOM(config, split='train'),
batch_size=self.batch_size,
shuffle=True,
num_workers=self.num_workers,
pin_memory=True)
self.val_loader = torch.utils.data.DataLoader(
DATASET_CUSTOM(config, split='val'),
batch_size=self.batch_size,
shuffle=False,
num_workers=self.num_workers,
pin_memory=True)
self.logger = Logger(config)
def main():
cfg = Config()
TVT, TMO = set_devices(cfg.sys_device_ids)
data_loader = get_data_loader(cfg)
spec_loss = SpectralCLusterLayer()
model = Model(cfg.vector_size, cfg.fix_weight)
model_w = DataParallel(model)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = cfg.lr)
modules_optims = [model, optimizer]
TMO(modules_optims)
may_set_mode(modules_optims, 'train')
for epoch in range(cfg.total_epoch):
epoch_done = False
step = 0
while not epoch_done:
step += 1
ims, _, labels, epoch_done= data_loader.next_batch()
ims_var = Variable(TVT(torch.from_numpy(ims).float()))
batch_size = ims_var.size()[0]
num_cluster = len(data_loader.ids)
labels_matrix = np.zeros([batch_size, num_cluster], dtype=int)
labels_matrix[range(batch_size), labels-1] = 1
labels_var = TVT(torch.from_numpy(labels_matrix).float())
optimizer.zero_grad()
feat = model_w(ims_var)
G = spec_loss.grad_F(feat, labels_var)
feat.backward(gradient=G)
optimizer.step()
objective_value = labels_var.size()[1] - torch.sum(torch.mm(spec_loss.pseudo_inverse(labels_var),feat) * torch.mm(spec_loss.pseudo_inverse(feat), labels_var).t())
print("epoch %d --- loss value= %f" % (epoch, objective_value))
print "Finished"
def __init__(self,
model,
regime=None,
criterion=None,
label_smoothing=0,
print_freq=10,
eval_freq=1000,
save_freq=1000,
grad_clip=None,
embedding_grad_clip=None,
limit_num_tokens=None,
save_info={},
save_path='.',
checkpoint_filename='checkpoint%s.pth.tar',
keep_checkpoints=5,
avg_loss_time=True,
distributed=False,
dist_local_rank=0,
dtype=torch.float,
device_ids=None,
device="cuda"):
super(Seq2SeqTrainer, self).__init__()
self.model = model
self.criterion = criterion or CrossEntropyLoss(
ignore_index=PAD, smooth_eps=label_smoothing, reduction='sum')
self.optimizer = OptimRegime(self.model.parameters(), regime=regime)
self.grad_clip = grad_clip
self.embedding_grad_clip = embedding_grad_clip
self.epoch = 0
self.training_steps = 0
self.save_info = save_info
self.device = device
self.dtype = dtype
self.limit_num_tokens = limit_num_tokens
self.print_freq = print_freq
self.eval_freq = eval_freq
self.perplexity = float('inf')
self.device_ids = device_ids
self.avg_loss_time = avg_loss_time
self.model_with_loss = AddLossModule(self.model, self.criterion)
self.distributed = distributed
if distributed:
self.model_with_loss = DistributedDataParallel(
self.model_with_loss,
device_ids=[dist_local_rank],
output_device=dist_local_rank)
else:
if isinstance(self.device_ids, tuple):
self.model_with_loss = DataParallel(
self.model_with_loss,
self.device_ids,
dim=0 if self.batch_first else 1)
self.save_path = save_path
self.save_freq = save_freq
self.checkpoint_filename = checkpoint_filename
self.keep_checkpoints = keep_checkpoints
results_file = os.path.join(save_path, 'results')
self.results = ResultsLog(results_file,
params=save_info.get('config', None))
def __init__(self, opt: Dict[str, Any]):
self.opt = opt
self.opt_train = self.opt['train']
self.opt_test = self.opt['test']
self.save_dir: str = opt['path']['models']
self.device = torch.device(
'cuda' if opt['gpu_ids'] is not None else 'cpu')
self.is_train = opt['is_train'] # training or not
self.type = opt['netG']['type']
self.net = select_network(opt).to(self.device)
self.net = DataParallel(self.net)
self.schedulers = []
self.log_dict = {}
self.metrics = {}
def load_model(filename, print_info=False):
use_cuda = torch.cuda.is_available()
chekpoint = torch.load(filename)
net = get_model(**chekpoint.get('script_args', {}))
net.load_state_dict(chekpoint['state_dict'])
if use_cuda:
net = DataParallel(net, device_ids=range(torch.cuda.device_count()))
return net
def test_attribute_evaluation():
evaluation, model_cfgs = evaluation_builder.build(attribute_evaluation_cfg)
with torch.no_grad():
for model_cfg in model_cfgs:
model = model_builder.build(model_cfg)
model = DataParallel(model)
score = evaluate(evaluation, model, delete=True)
print(score)
def _build_model(self):
""" Build image retrieval network and load pre-trained weights.
"""
model = nn.Module()
# Assume a VGG-16 backbone
encoder = models.vgg16(pretrained=False)
layers = list(encoder.features.children())[:-2]
encoder = nn.Sequential(*layers)
model.add_module('encoder', encoder)
# Assume a NetVLAD pooling layer
net_vlad = NetVLAD(num_clusters=self._num_clusters,
dim=self._encoder_dim)
model.add_module('pool', net_vlad)
# For parallel training
gpu_count = torch.cuda.device_count()
if gpu_count > 1:
model.encoder = DataParallel(model.encoder)
model.pool = DataParallel(model.pool)
# Load weights
checkpoint = torch.load(
self._checkpoint_path,
map_location=lambda storage, loc: storage)['state_dict']
# If model was not trained in parallel, adapt the keys
if 'module' not in list(checkpoint.keys())[0] and False:
checkpoint = OrderedDict(
(k.replace('encoder', 'encoder.module'), v)
for k, v in checkpoint.items())
checkpoint = OrderedDict((k.replace('pool', 'pool.module'), v)
for k, v in checkpoint.items())
elif gpu_count <= 1:
checkpoint = OrderedDict(
(k.replace('encoder.module', 'encoder'), v)
for k, v in checkpoint.items())
checkpoint = OrderedDict((k.replace('pool.module', 'pool'), v)
for k, v in checkpoint.items())
model.load_state_dict(checkpoint)
if self._device != -1:
model = model.to(self._device)
model.eval()
return model
def get_learner(
PATH,
path_x,
path_y,
save_dir,
path_stats,
bs,
sz,
gpu_start,
arch='DlinkNet34',
world_size=1,
test_size=0.2,
num_workers=16,
resume=False,
use_clr=None,
cycle_len=1,
lr=0.01,
momentum=0.9,
wd=1e-4,
epochs=1, # debug
loss_scale=1,
load_model=None,
path_test_x=None,
path_test_y=None):
global md, print_freq, _use_clr, _cycle_len, _rank
global learner, denorm
_use_clr, _cycle_len, _save_dir = use_clr, cycle_len, save_dir
torch.cuda.set_device(gpu_start)
device_ids = range(gpu_start, gpu_start + world_size)
stats = np.load(path_stats)
md, denorm = get_loader(PATH,
path_x,
path_y,
stats,
bs,
sz,
classes=7,
test_size=test_size,
num_workers=num_workers,
path_test_x=path_test_x,
path_test_y=path_test_y)
# This is important for speed
cudnn.benchmark = True
distributed = world_size > 1
model = models.__dict__[arch](num_classes=7)
if distributed: model = DataParallel(model, device_ids)
learner = Learner.from_model_data(model, md)
learner.crit = F.cross_entropy
if load_model is not None:
sd = torch.load(load_model, map_location=lambda storage, loc: storage)
learner.model.load_state_dict(sd)
# Full size
update_model_dir(learner, save_dir)
sargs = save_args('first_run', save_dir)
def __init__(self, args):
super(GANLoss, self).__init__()
# By default, discriminator will be updated every step
if args.discriminator == 'discriminator_vgg_128':
self.netD = Discriminator_VGG_128(in_nc=3, nf=64)
else:
raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(args.discriminator))
if args.pretrained_netD is not None:
self.netD.load_state_dict(torch.load(args.pretrained_netD))
if not args.cpu:
self.netD = self.netD.to('cuda')
if args.n_GPUs > 1:
self.netD = DataParallel(self.netD)
self.save_D_every = args.save_D_every
if args.save_D_path == '...':
self.save_D_path = "../experiments/{}/model/".format(args.name)
else:
self.save_D_path = args.save_D_path
# Loss Type
self.gan_type = args.gan_type
if self.gan_type == 'gan' or self.gan_type == 'ragan':
self.loss = nn.BCEWithLogitsLoss()
elif self.gan_type == 'lsgan':
self.loss = nn.MSELoss()
elif self.gan_type == 'wgan-gp':
def wgan_loss(input, target):
# target is boolean
return -1 * input.mean() if target else input.mean()
self.loss = wgan_loss
else:
raise NotImplementedError('GAN type [{:s}] is not found'.format(self.gan_type))
# Optimizer
self.optimizer_D = torch.optim.Adam(
params = self.netD.parameters(),
lr = args.lr_D,
betas = (args.beta1_D, args.beta2_D),
weight_decay = args.weight_decay_D
)
def __init__(self,
model,
images_path,
labels_path,
patient_ids,
sample_shape,
checkpoint_restore,
inference_dir,
use_gpu=False,
gpu_ids=None):
# model settings
self.model = model
# data settings
assert len(images_path) == len(labels_path)
self.images_path = images_path
self.labels_path = labels_path
self.patient_ids = patient_ids
self.length = len(images_path)
self.sample_shape = sample_shape
self.inference_dir = inference_dir
# gpu settings
self.use_gpu = use_gpu
if use_gpu and torch.cuda.device_count() > 0:
self.model.cuda()
if gpu_ids is not None:
if len(gpu_ids) > 1:
self.multi_gpu = True
self.model = DataParallel(model, gpu_ids)
else:
self.multi_gpu = False
else:
if torch.cuda.device_count() > 1:
self.multi_gpu = True
self.model = DataParallel(model)
else:
self.multi_gpu = False
else:
self.multi_gpu = False
self.model = self.model.cpu()
self.model.load_state_dict(torch.load(checkpoint_restore))
def main():
TVT, TMO = set_devices((-1,))
model = Model(local_conv_out_channels=128,
num_classes=751)
# Model wrapper
model_w = DataParallel(model)
map_location = (lambda storage, loc: storage)
sd = torch.load('ckpt.pth', map_location=map_location)
load_state_dict(model, sd)
print('Loaded model weights')
IMAGES_PATH = "galery/"
files = os.listdir(IMAGES_PATH)
files.sort();
data = []
formats = ["png", "jpg"]
curr_id = 0
ids = {}
for image_name in files:
if image_name.split('.')[1] in formats and int(image_name.split('_')[0]) > 0:
id = image_name.split('_')[0]
'''
if id not in ids:
ids[id] = 0
else:
if ids[id] > 20:
continue
ids[id]+=1
'''
im = PIL.Image.open(IMAGES_PATH + image_name)
im = im.resize((128,256))
#im.save('tmp.jpg')
I = np.asarray(im)
I = I - np.mean(I)
I = I/np.std(I)
I= np.expand_dims(I.transpose(2, 0, 1),axis = 0)
ims = Variable(TVT(torch.from_numpy(I).float()))
global_feat, local_feat = model(ims)[:2]
global_feat = np.vstack(global_feat.data.cpu().numpy())
local_feat = local_feat.data.cpu().numpy()
#print(features)
print("processing file: ", image_name)
data.append({"id" : id, "features" : (global_feat, local_feat),"path" : IMAGES_PATH + image_name})
filename = 'data'
outfile = open(filename,'wb')
pickle.dump(data,outfile)
outfile.close()
