def main(args=None):
args, no_config_file = get_args(args)
setup_log(args.verbosity)
logger = logging.getLogger(__name__)
try:
# set random seeds for reproducibility
torch.manual_seed(args.seed)
np.random.seed(args.seed)
# since prediction only uses one gpu (at most), make the batch size small enough to fit
if args.n_gpus > 1: args.batch_size = args.batch_size // args.n_gpus
learner = Learner.predict_setup(args)
# determine how many samples we will use in prediction
nsyn = args.monte_carlo or 1
# get relevant prediction directories and determine extension
predict_dir = args.predict_dir or args.valid_source_dir
output_dir = args.predict_out or os.getcwd() + '/syn_'
ext = determine_ext(predict_dir[0])
# setup and start prediction loop
axis = args.sample_axis or 0
if axis < 0 or axis > 2 and not isinstance(axis, int):
raise ValueError(
'sample_axis must be an integer between 0 and 2 inclusive')
n_imgs = len(glob_imgs(predict_dir[0], ))
if n_imgs == 0:
raise SynthNNError(
'Prediction directory does not contain valid images.')
if any([len(glob_imgs(pd, ext)) != n_imgs for pd in predict_dir]):
raise SynthNNError(
'Number of images in prediction directories must have an equal number of images in each '
'directory (e.g., so that img_t1_1 aligns with img_t2_1 etc. for multimodal synth)'
)
predict_fns = zip(*[glob_imgs(pd, ext) for pd in predict_dir])
if args.net3d and args.patch_size > 0 and args.calc_var:
raise SynthNNError(
'Patch-based 3D variance calculation not currently supported.')
for k, fn in enumerate(predict_fns):
_, base, ext = split_filename(fn[0])
logger.info(
f'Starting synthesis of image: {base} ({k+1}/{n_imgs})')
out_imgs = learner.predict(fn, nsyn, args.temperature_map,
args.calc_var)
for i, oin in enumerate(out_imgs):
out_fn = output_dir + f'{k}_{i}{ext}'
if hasattr(oin, 'to_filename'):
oin.to_filename(out_fn)
else:
oin.save(out_fn)
logger.info(f'Finished synthesis. Saved as: {out_fn}')
return 0
except Exception as e:
logger.exception(e)
return 1
def setUp(self):
wd = os.path.dirname(os.path.abspath(__file__))
self.nii_dir = os.path.join(wd, 'test_data', 'nii')
self.mask_dir = os.path.join(wd, 'test_data', 'masks')
self.tif_dir = os.path.join(wd, 'test_data', 'tif')
self.png_dir = os.path.join(wd, 'test_data', 'png')
self.out_dir = tempfile.mkdtemp()
os.mkdir(os.path.join(self.out_dir, 'models'))
self.train_dir = os.path.join(self.out_dir, 'imgs')
os.mkdir(self.train_dir)
os.mkdir(os.path.join(self.train_dir, 'mask'))
os.mkdir(os.path.join(self.train_dir, 'tif'))
os.mkdir(os.path.join(self.train_dir, 'png'))
nii = glob_imgs(self.nii_dir)[0]
msk = glob_imgs(self.mask_dir)[0]
tif = os.path.join(self.tif_dir, 'test.tif')
png = os.path.join(self.png_dir, 'test.png')
path, base, ext = split_filename(nii)
for i in range(8):
shutil.copy(nii, os.path.join(self.train_dir, base + str(i) + ext))
shutil.copy(
msk, os.path.join(self.train_dir, 'mask', base + str(i) + ext))
shutil.copy(
tif, os.path.join(self.train_dir, 'tif',
base + str(i) + '.tif'))
shutil.copy(
png, os.path.join(self.train_dir, 'png',
base + str(i) + '.png'))
self.train_args = f'-s {self.train_dir} -t {self.train_dir}'.split()
self.predict_args = f'-s {self.train_dir} -o {self.out_dir}/test'.split(
)
self.jsonfn = f'{self.out_dir}/test.json'
def fit(self,
n_epochs,
clip: float = None,
checkpoint: int = None,
trained_model: str = None):
""" training loop for neural network """
self.model.train()
use_valid = self.valid_loader is not None
train_losses, valid_losses = [], []
for t in range(1, n_epochs + 1):
# training
t_losses = []
if use_valid: self.model.train(True)
for src, tgt in self.train_loader:
src, tgt = src.to(self.device), tgt.to(self.device)
self.optimizer.zero_grad()
out = self.model(src)
loss = self._criterion(out, tgt)
t_losses.append(loss.item())
if hasattr(self, 'amp_handle'):
with self.amp_handle.scale_loss(
loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
if clip is not None:
nn.utils.clip_grad_norm_(self.model.parameters(), clip)
self.optimizer.step()
train_losses.append(t_losses)
if hasattr(self, 'scheduler'): self.scheduler.step()
if checkpoint is not None:
if t % checkpoint == 0:
path, base, ext = split_filename(trained_model)
fn = os.path.join(path, base + f'_chk_{t}' + ext)
self.save(fn, t)
# validation
v_losses = []
if use_valid: self.model.train(False)
with torch.set_grad_enabled(False):
for src, tgt in self.valid_loader:
src, tgt = src.to(self.device), tgt.to(self.device)
out = self.model(src)
loss = self._criterion(out, tgt)
v_losses.append(loss.item())
valid_losses.append(v_losses)
if np.any(np.isnan(t_losses)):
raise SynthNNError(
'NaN in training loss, cannot recover. Exiting.')
if logger is not None:
log = f'Epoch: {t} - Training Loss: {np.mean(t_losses):.2e}'
if use_valid:
log += f', Validation Loss: {np.mean(v_losses):.2e}'
if hasattr(self, 'scheduler'):
log += f', LR: {self.scheduler.get_lr()[0]:.2e}'
logger.info(log)
return train_losses, valid_losses
def fit(self, n_epochs, clip:float=None, checkpoint:int=None, trained_model:str=None):
""" training loop for neural network """
self.model.train()
use_valid = self.valid_loader is not None
use_scheduler = hasattr(self, 'scheduler')
use_restarts = self.config.lr_scheduler == 'cosinerestarts'
fp16 = hasattr(self, 'amp_handle')
train_losses, valid_losses = [], []
n_batches = len(self.train_loader)
for t in range(1, n_epochs + 1):
# training
t_losses = []
if use_valid: self.model.train(True)
for i, (src, tgt) in enumerate(self.train_loader):
src, tgt = src.to(self.device), tgt.to(self.device)
self.optimizer.zero_grad()
out = self.model(src)
loss = self._criterion(out, tgt)
t_losses.append(loss.item())
if fp16:
with self.amp_handle.scale_loss(loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
if clip is not None: nn.utils.clip_grad_norm_(self.model.parameters(), clip)
self.optimizer.step()
if use_scheduler: self.scheduler.step(((t - 1) + (i / n_batches)) if use_restarts else None)
del loss # save memory by removing ref to gradient tree
train_losses.append(t_losses)
if checkpoint is not None:
if t % checkpoint == 0:
path, base, ext = split_filename(trained_model)
fn = os.path.join(path, base + f'_chk_{t}' + ext)
self.save(fn, t)
# validation
v_losses = []
if use_valid: self.model.train(False)
with torch.no_grad():
for src, tgt in self.valid_loader:
src, tgt = src.to(self.device), tgt.to(self.device)
out = self.model(src)
loss = self._criterion(out, tgt)
v_losses.append(loss.item())
valid_losses.append(v_losses)
if not np.all(np.isfinite(t_losses)): raise SynthtorchError('NaN or Inf in training loss, cannot recover. Exiting.')
if logger is not None:
log = f'Epoch: {t} - Training Loss: {np.mean(t_losses):.2e}'
if use_valid: log += f', Validation Loss: {np.mean(v_losses):.2e}'
if use_scheduler: log += f', LR: {self.scheduler.get_lr()[0]:.2e}'
logger.info(log)
self.record = Record(train_losses, valid_losses)
def fit(self, n_epochs, clip:float=None, checkpoint:int=None, trained_model:str=None):
""" training loop for neural network """
self.model.train()
use_tb = self.config.tensorboard and SummaryWriter is not None
if use_tb: writer = SummaryWriter()
use_valid = self.valid_loader is not None
use_scheduler = hasattr(self, 'scheduler')
use_restarts = self.config.lr_scheduler == 'cosinerestarts'
train_losses, valid_losses = [], []
n_batches = len(self.train_loader)
for t in range(1, n_epochs + 1):
# training
t_losses = []
if use_valid: self.model.train(True)
for i, (src, tgt) in enumerate(self.train_loader):
src, tgt = src.to(self.device), tgt.to(self.device)
self.optimizer.zero_grad()
out = self.model(src)
loss = self._criterion(out, tgt)
t_losses.append(loss.item())
if self.use_fp16:
with amp.scale_loss(loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
if clip is not None: nn.utils.clip_grad_norm_(self.model.parameters(), clip)
self.optimizer.step()
if use_scheduler: self.scheduler.step(((t-1)+(i/n_batches)) if use_restarts else None)
if use_tb:
if i % 20 == 0: writer.add_scalar('Loss/train', loss.item(), ((t-1)*n_batches)+i)
del loss # save memory by removing ref to gradient tree
train_losses.append(t_losses)
if checkpoint is not None:
if t % checkpoint == 0:
path, base, ext = split_filename(trained_model)
fn = os.path.join(path, base + f'_chk_{t}' + ext)
self.save(fn, t)
# validation
v_losses = []
if use_valid:
self.model.train(False)
with torch.no_grad():
for i, (src, tgt) in enumerate(self.valid_loader):
src, tgt = src.to(self.device), tgt.to(self.device)
out = self.model(src)
loss = self._criterion(out, tgt)
if use_tb:
if i % 20 == 0: writer.add_scalar('Loss/valid', loss.item(), ((t-1)*n_batches)+i)
do_plot = i == 0 and ((t - 1) % 5) == 0
if do_plot and self.model.dim == 2:
writer.add_images('source', src[:8], t, dataformats='NCHW')
outimg = out[0][:8] if isinstance(out, tuple) else out[:8]
if self.config.color: outimg = torch.round(outimg)
writer.add_images('target', outimg, t, dataformats='NCHW')
if do_plot: self._histogram_weights(writer, t)
v_losses.append(loss.item())
valid_losses.append(v_losses)
if not np.all(np.isfinite(t_losses)): raise SynthtorchError('NaN or Inf in training loss, cannot recover. Exiting.')
if logger is not None:
log = f'Epoch: {t} - Training Loss: {np.mean(t_losses):.2e}'
if use_valid: log += f', Validation Loss: {np.mean(v_losses):.2e}'
if use_scheduler: log += f', LR: {self.scheduler.get_lr()[0]:.2e}'
logger.info(log)
self.record = Record(train_losses, valid_losses)
if use_tb: writer.close()