def load_data2():
### https://github.com/HasnainRaz/SemSegPipeline
data_dir = base_dir + '/ultrasound-nerve-segmentation/train_jpeg'
image_paths = [
os.path.join(data_dir, x) for x in os.listdir(data_dir)
if x.endswith('.tif') and 'mask' not in x
]
mask_paths = [
os.path.join(data_dir, x) for x in os.listdir(data_dir)
if x.endswith('.tif') and 'mask' in x
]
# print(image_paths)
# print(mask_paths)
dataset = DataLoader(image_paths=image_paths,
mask_paths=mask_paths,
image_size=[256, 256],
crop_percent=1.0,
channels=[3, 3],
seed=47)
dataset = dataset.data_batch(batch_size=BATCH_SIZE,
augment=False,
shuffle=True)
print(dataset)
def load_data(self):
"""Loads and Preprocess data """
LOG.info(f'Loading {self.config.data.path} dataset...')
self.dataset, self.info = DataLoader().load_data(self.config.data)
self.train_dataset, self.test_dataset = DataLoader.preprocess_data(
self.dataset, self.batch_size, self.buffer_size, self.image_size)
self._set_training_parameters()
def __init__(self, args):
self.args = args
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if 'audio' in args.generate:
print(f"----------INPUT SIZE IS {args.cond_size+args.gen_size}")
self.model = CondRNN(args.cond_size+args.gen_size, args.hidden_size,
args.mulaw_channels, args.n_layers, self.device,
args.lr,paramonly=False,onehot=args.onehot,
ntype=args.net,plstm=args.plstm)
else:
self.model = CondRNN(args.cond_size+args.gen_size, args.hidden_size,
args.gen_size, args.n_layers, self.device,
args.lr,paramonly=True,onehot=args.onehot,
ntype=args.net,plstm=args.plstm)
self.data_loader = DataLoader(args.data_dir, args.sample_rate, args.seq_len, args.stride,
paramdir=args.param_dir, prop=args.prop, generate=args.generate,
mulaw_channels=args.mulaw_channels,
batch_size=args.batch_size,
onehot=args.onehot)
self.epoch_size = len(self.data_loader)
if args.model_dir: #to resume from checkpoint
self.model.load(args.model_dir, args.step)
print("Resuming training from step",args.step)
def __init__(self, args):
self.args = args
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
print("using", self.device)
#if torch.cuda.is_available():
# torch.cuda.empty_cache()
if 'audio' in args.generate:
self.model = CondRNN(args.cond_size + args.gen_size,
args.hidden_size,
args.mulaw_channels,
args.n_layers,
self.device,
paramonly=False,
onehot=args.onehot,
ntype=args.net,
plstm=args.plstm)
else:
self.model = CondRNN(args.cond_size + args.gen_size,
args.hidden_size,
args.gen_size,
args.n_layers,
self.device,
paramonly=True,
onehot=args.onehot,
ntype=args.net,
plstm=args.plstm)
self.model.load(args.model_dir, args.step)
self.prime_length = self.args.seq_len - self.args.length
if args.seed is not None:
self._get_seed_from_audio(self.args.seed)
elif args.seed is None and args.data_dir is None:
assert args.rand_prime, "provide either a seed file or directory for priming/generation! Or use rand_prime."
assert self.args.paramvect != 'self', "provide either a seed file or directory for 'self' generation!"
else:
if self.args.paramvect == 'self':
load_length = self.args.seq_len
else:
load_length = self.prime_length
self.data_loader = DataLoader(args.data_dir,
args.sample_rate,
load_length,
args.stride,
paramdir=args.param_dir,
prop=args.prop,
generate=args.generate,
mulaw_channels=args.mulaw_channels,
batch_size=args.batch_size,
onehot=args.onehot)
if args.external_array is not None:
self.param_array = np.load(args.external_array)
def train():
data_loader = DataLoader(data_dir='/media/DATASET/wangning/data/建筑材质',
image_size=IMAGE_SIZE,
batch_size=64)
inputs, classes = next(iter(data_loader.load_data()))
out = torchvision.utils.make_grid(inputs)
data_loader.show_image(
out, title=[data_loader.data_classes[c] for c in classes])
model = fine_tune_model()
criterion = nn.CrossEntropyLoss()
ignored_params = list(map(id, model.fc.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer_ft = optim.SGD([{
'params': base_params
}, {
'params': model.fc.parameters(),
'lr': 1e-2
}],
lr=1e-3,
momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft,
step_size=7,
gamma=0.1)
try:
model = train_model(data_loader,
model,
criterion,
optimizer_ft,
exp_lr_scheduler,
num_epochs=50)
save_torch_model(model, MODEL_SAVE_FILE)
except KeyboardInterrupt:
print('manually interrupt, try saving model for now...')
save_torch_model(model, MODEL_SAVE_FILE)
print('model saved.')
def evaluate(self):
"""Predicts resuts for the test dataset"""
predictions = []
LOG.info('Predicting segmentation map for test dataset')
for im in self.test_dataset.as_numpy_iterator():
DataLoader().validate_schema(im[0])
break
for image, mask in self.test_dataset:
tf.print(image)
# LOG.debug(f'Predicting segmentation map {image}')
predictions.append(self.model.predict(image))
return predictions
def _get_seed_from_audio(self, filepath):
if self.args.paramvect == 'self':
load_length = self.args.seq_len
else:
load_length = self.prime_length
self.data_loader = DataLoader(filepath,
self.args.sample_rate,
load_length,
self.args.stride,
paramdir=self.args.param_dir,
prop=self.args.prop,
generate=self.args.generate,
mulaw_channels=self.args.mulaw_channels,
batch_size=self.args.batch_size,
shuffle=False,
onehot=self.args.onehot,
is_seeded=True,
load_start=self.args.seed_start)
def load_data(self):
"""Loads and Preprocess data"""
LOG.info(f'loading {self.config.data.path} dataset')
self.dataset, self.info = DataLoader().load_data(self.config.data)
self._preprocess_data()
def load_data(self):
"""Loads and Preprocess data """
self.dataset = DataLoader().load_data(self.config)
self._preprocess_data()
class UNet(BaseModel):
"""Unet Model Class"""
def __init__(self, config):
super().__init__(config)
# self.base_model = tf.keras.applications.MobileNetV2(input_shape=self.config.model.input, include_top=False)
#https: // www.tensorflow.org / api_docs / python / tf / keras / applications?hl = de
self.model = None
self.output_channels = self.config.model.output
self.dataset = None
self.info = None
self.batch_size = self.config.train.batch_size
#self.buffer_size = self.config.train.buffer_size
self.epoches = self.config.train.epoches
#self.val_subsplits = self.config.train.val_subsplits
self.validation_steps = 0
self.train_length = 0
self.steps_per_epoch = 0
self.image_size = self.config.data.image_size
self.train_dataset = []
self.test_dataset = []
def load_data(self):
"""Loads and Preprocess data """
self.dataset = DataLoader().load_data(self.config)
self._preprocess_data()
def _preprocess_data(self):
""" Splits into training and test and set training parameters"""
self.train_dataset = self.dataset['generated_training_batches']
self.test_dataset = self.dataset['generated_test_batches']
#self._set_training_parameters()
# def _set_training_parameters(self):
# """Sets training parameters"""
#self.train_length = self.info.splits['train'].num_examples
#self.steps_per_epoch = self.train_length // self.batch_size
#self.validation_steps = self.info.splits['test'].num_examples // self.batch_size // self.val_subsplits
# def _normalize(self, input_image, input_mask):
# !!! normalization by scaling happends in the dataloader.py by ImageDataGenerator !!!
# """ Normalise input image
# Args:
# input_image (tf.image): The input image
# input_mask (int): The image mask
#
# Returns:
# input_image (tf.image): The normalized input image
# input_mask (int): The new image mask
# """
# input_image = tf.cast(input_image, tf.float32) / 255.0
# input_mask -= 1
# return input_image, input_mask
# @tf.function
# def _load_image_train(self, datapoint):
# """ Loads and preprocess a single training image """
# input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size))
# input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size))
#
# if tf.random.uniform(()) > 0.5:
# input_image = tf.image.flip_left_right(input_image)
# input_mask = tf.image.flip_left_right(input_mask)
#
# input_image, input_mask = self._normalize(input_image, input_mask)
#
# return input_image, input_mask
# def _load_image_test(self, datapoint):
# """ Loads and preprocess a single test images"""
#
# input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size))
# input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size))
#
# input_image, input_mask = self._normalize(input_image, input_mask)
#
# return input_image, input_mask
def build(self):
""" Builds the Keras model based """
layer_names = [
'block_1_expand_relu', # 64x64
'block_3_expand_relu', # 32x32
'block_6_expand_relu', # 16x16
'block_13_expand_relu', # 8x8
'block_16_project', # 4x4
]
layers = [self.base_model.get_layer(name).output for name in layer_names]
# Create the feature extraction model
# down_stack = tf.keras.Model(inputs=self.base_model.input, outputs=layers)
# down_stack.trainable = False
up_stack = [
pix2pix.upsample(self.config.model.up_stack.layer_1, self.config.model.up_stack.kernels), # 4x4 -> 8x8
pix2pix.upsample(self.config.model.up_stack.layer_2, self.config.model.up_stack.kernels), # 8x8 -> 16x16
pix2pix.upsample(self.config.model.up_stack.layer_3, self.config.model.up_stack.kernels), # 16x16 -> 32x32
pix2pix.upsample(self.config.model.up_stack.layer_4, self.config.model.up_stack.kernels), # 32x32 -> 64x64
]
inputs = tf.keras.layers.Input(shape=self.config.model.input)
x = inputs
# Downsampling through the model
# skips = down_stack(x)
# x = skips[-1]
# skips = reversed(skips[:-1])
# Upsampling and establishing the skip connections
# for up, skip in zip(up_stack, skips):
# x = up(x)
# concat = tf.keras.layers.Concatenate()
# x = concat([x, skip])
# This is the last layer of the model
last = tf.keras.layers.Conv2DTranspose(
self.output_channels, self.config.model.up_stack.kernels, strides=2,
padding='same') # 64x64 -> 128x128
x = last(x)
self.model = tf.keras.Model(inputs=inputs, outputs=x)
def train(self):
"""Compiles and trains the model"""
self.model.compile(optimizer=self.config.train.optimizer.type,
loss=tf.keras.losses.BinaryCrossentropy(),
metrics=self.config.train.metrics)
early_stop = EarlyStopping(monitor=self.config.train.EarlyStopping.monitor,
mode=self.config.train.EarlyStopping.mode,
verbose=1,
patience=self.config.train.EarlyStopping.self.config.train.EarlyStopping.patience)
model_history = self.model.fit_generator(self.train_dataset,
epochs=self.epoches,
steps_per_epoch=self.steps_per_epoch,
validation_steps=self.validation_steps,
callbacks=[early_stop],
validation_data=self.test_dataset)
return model_history.history.history
def evaluate(self):
"""Predicts resuts for the test dataset"""
predictions = []
for image, mask in self.dataset.take(1):
predictions.append(self.model.predict(image))
return predictions
def distributed_train(self):
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
with mirrored_strategy.scope():
self.model = tf.keras.Model(inputs=inputs, outputs=x)
self.model.compile(...)
self.model.fit(...)
os.environ["TF_CONFIG"] = json.dumps(
{
"cluster":{
"worker": ["host1:port", "host2:port", "host3:port"]
},
"task":{
"type": "worker",
"index": 1
}
}
)
multi_worker_mirrored_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
with multi_worker_mirrored_strategy.scope():
self.model = tf.keras.Model(inputs=inputs, outputs=x)
self.model.compile(...)
self.model.fit(...)
parameter_server_strategy = tf.distribute.experimental.ParameterServerStrategy()
os.environ["TF_CONFIG"] = json.dumps(
{
"cluster": {
"worker": ["host1:port", "host2:port", "host3:port"],
"ps": ["host4:port", "host5:port"]
},
"task": {
"type": "worker",
"index": 1
}
}
)
class UNet(BaseModel):
"""Unet Model Class"""
def __init__(self, config):
super().__init__(config)
self.base_model = tf.keras.applications.MobileNetV2(input_shape=self.config.model.input, include_top=False)
self.model = None
self.output_channels = self.config.model.output
self.dataset = None
self.info = None
self.batch_size = self.config.train.batch_size
self.buffer_size = self.config.train.buffer_size
self.epoches = self.config.train.epoches
self.val_subsplits = self.config.train.val_subsplits
self.validation_steps = 0
self.train_length = 0
self.steps_per_epoch = 0
self.image_size = self.config.data.image_size
self.train_dataset = []
self.test_dataset = []
def load_data(self):
"""Loads and Preprocess data """
self.dataset, self.info = DataLoader().load_data(self.config.data)
self._preprocess_data()
def _preprocess_data(self):
""" Splits into training and test and set training parameters"""
train = self.dataset['train'].map(self._load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)
test = self.dataset['test'].map(self._load_image_test)
self.train_dataset = train.cache().shuffle(self.buffer_size).batch(self.batch_size).repeat()
self.train_dataset = self.train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
self.test_dataset = test.batch(self.batch_size)
self._set_training_parameters()
def _set_training_parameters(self):
"""Sets training parameters"""
self.train_length = self.info.splits['train'].num_examples
self.steps_per_epoch = self.train_length // self.batch_size
self.validation_steps = self.info.splits['test'].num_examples // self.batch_size // self.val_subsplits
def _normalize(self, input_image, input_mask):
""" Normalise input image
Args:
input_image (tf.image): The input image
input_mask (int): The image mask
Returns:
input_image (tf.image): The normalized input image
input_mask (int): The new image mask
"""
input_image = tf.cast(input_image, tf.float32) / 255.0
input_mask -= 1
return input_image, input_mask
@tf.function
def _load_image_train(self, datapoint):
""" Loads and preprocess a single training image """
input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size))
input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size))
if tf.random.uniform(()) > 0.5:
input_image = tf.image.flip_left_right(input_image)
input_mask = tf.image.flip_left_right(input_mask)
input_image, input_mask = self._normalize(input_image, input_mask)
return input_image, input_mask
def _load_image_test(self, datapoint):
""" Loads and preprocess a single test images"""
input_image = tf.image.resize(datapoint['image'], (self.image_size, self.image_size))
input_mask = tf.image.resize(datapoint['segmentation_mask'], (self.image_size, self.image_size))
input_image, input_mask = self._normalize(input_image, input_mask)
return input_image, input_mask
def build(self):
""" Builds the Keras model based """
layer_names = [
'block_1_expand_relu', # 64x64
'block_3_expand_relu', # 32x32
'block_6_expand_relu', # 16x16
'block_13_expand_relu', # 8x8
'block_16_project', # 4x4
]
layers = [self.base_model.get_layer(name).output for name in layer_names]
# Create the feature extraction model
down_stack = tf.keras.Model(inputs=self.base_model.input, outputs=layers)
down_stack.trainable = False
up_stack = [
pix2pix.upsample(self.config.model.up_stack.layer_1, self.config.model.up_stack.kernels), # 4x4 -> 8x8
pix2pix.upsample(self.config.model.up_stack.layer_2, self.config.model.up_stack.kernels), # 8x8 -> 16x16
pix2pix.upsample(self.config.model.up_stack.layer_3, self.config.model.up_stack.kernels), # 16x16 -> 32x32
pix2pix.upsample(self.config.model.up_stack.layer_4, self.config.model.up_stack.kernels), # 32x32 -> 64x64
]
inputs = tf.keras.layers.Input(shape=self.config.model.input)
x = inputs
# Downsampling through the model
skips = down_stack(x)
x = skips[-1]
skips = reversed(skips[:-1])
# Upsampling and establishing the skip connections
for up, skip in zip(up_stack, skips):
x = up(x)
concat = tf.keras.layers.Concatenate()
x = concat([x, skip])
# This is the last layer of the model
last = tf.keras.layers.Conv2DTranspose(
self.output_channels, self.config.model.up_stack.kernels, strides=2,
padding='same') # 64x64 -> 128x128
x = last(x)
self.model = tf.keras.Model(inputs=inputs, outputs=x)
def train(self):
"""Compiles and trains the model"""
self.model.compile(optimizer=self.config.train.optimizer.type,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=self.config.train.metrics)
model_history = self.model.fit(self.train_dataset, epochs=self.epoches,
steps_per_epoch=self.steps_per_epoch,
validation_steps=self.validation_steps,
validation_data=self.test_dataset)
return model_history.history['loss'], model_history.history['val_loss']
def evaluate(self):
"""Predicts resuts for the test dataset"""
predictions = []
for image, mask in self.dataset.take(1):
predictions.append(self.model.predict(image))
return predictions