From 9bec563f7508344f41bc191dc348f945597bd37b Mon Sep 17 00:00:00 2001 From: dirk Date: Tue, 28 Mar 2023 10:17:06 +0200 Subject: [PATCH 1/7] Add .gitignore file --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ab06921 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +*~ +.venv +__pycache__ From 03b098ea4b4b69b8df1fb1c77feb7cae7c53139a Mon Sep 17 00:00:00 2001 From: dirk Date: Tue, 28 Mar 2023 14:15:44 +0200 Subject: [PATCH 2/7] Add swp-files to .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index ab06921..db49ef0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ *~ +*.swp .venv __pycache__ From 906fc0354c3a91c35b7e9e85465c233703838bbf Mon Sep 17 00:00:00 2001 From: dirk Date: Tue, 28 Mar 2023 14:16:24 +0200 Subject: [PATCH 3/7] Separate out 2D layers --- blob.py | 56 ++++ caffemodel2pytorch.py | 765 ++++++++++++++++++++++-------------------- layers2d.py | 291 ++++++++++++++++ test.py | 41 +++ utils.py | 72 ++++ 5 files changed, 860 insertions(+), 365 deletions(-) create mode 100644 blob.py create mode 100644 layers2d.py create mode 100644 test.py create mode 100644 utils.py diff --git a/blob.py b/blob.py new file mode 100644 index 0000000..96b9dd6 --- /dev/null +++ b/blob.py @@ -0,0 +1,56 @@ +from functools import reduce + + +class Blob(object): + AssignmentAdapter = type( + "", + (object,), + dict( + shape=property(lambda self: self.contents.shape), + __setitem__=lambda self, indices, values: setattr(self, "contents", values), + ), + ) + + def __init__(self, data=None, diff=None, numpy=False): + self.data_ = data if data is not None else Blob.AssignmentAdapter() + self.diff_ = diff if diff is not None else Blob.AssignmentAdapter() + self.shape_ = None + self.numpy = numpy + + def reshape(self, *args): + self.shape_ = args + + def count(self, *axis): + return reduce(lambda x, y: x * y, self.shape_[slice(*(axis + [-1])[:2])]) + + @property + def data(self): + if self.numpy and torch.is_tensor(self.data_): + self.data_ = self.data_.detach().cpu().numpy() + return self.data_ + + @property + def diff(self): + if self.numpy and torch.is_tensor(self.diff_): + self.diff_ = self.diff_.detach().cpu().numpy() + return self.diff_ + + @property + def shape(self): + return self.shape_ if self.shape_ is not None else self.data_.shape + + @property + def num(self): + return self.shape[0] + + @property + def channels(self): + return self.shape[1] + + @property + def height(self): + return self.shape[2] + + @property + def width(self): + return self.shape[3] diff --git a/caffemodel2pytorch.py b/caffemodel2pytorch.py index 6e705b4..4a8e2e2 100755 --- a/caffemodel2pytorch.py +++ b/caffemodel2pytorch.py @@ -9,17 +9,21 @@ import torch.nn as nn import torch.nn.functional as F from functools import reduce - -try: - from urllib.request import urlopen -except: - from urllib2 import urlopen # Python 2 support. +from urllib.request import urlopen import google.protobuf.descriptor import google.protobuf.descriptor_pool import google.protobuf.symbol_database import google.protobuf.text_format -from google.protobuf.descriptor import FieldDescriptor as FD +import google.protobuf.json_format + +import ssl + +from layers2d import modules2d +from blob import Blob +from utils import to_dict, convert_to_gpu_if_enabled + +ssl._create_default_https_context = ssl._create_unverified_context TRAIN = 0 @@ -27,373 +31,404 @@ caffe_pb2 = None -def initialize(caffe_proto = 'https://raw.githubusercontent.com/BVLC/caffe/master/src/caffe/proto/caffe.proto', codegen_dir = tempfile.mkdtemp(), shadow_caffe = True): - global caffe_pb2 - if caffe_pb2 is None: - local_caffe_proto = os.path.join(codegen_dir, os.path.basename(caffe_proto)) - with open(local_caffe_proto, 'w') as f: - mybytes = urlopen(caffe_proto).read() - mystr = mybytes.decode('ascii', 'ignore') - f.write(mystr) - #f.write((urlopen if 'http' in caffe_proto else open)(caffe_proto).read()) - subprocess.check_call(['protoc', '--proto_path', os.path.dirname(local_caffe_proto), '--python_out', codegen_dir, local_caffe_proto]) - sys.path.insert(0, codegen_dir) - old_symdb = google.protobuf.symbol_database._DEFAULT - google.protobuf.symbol_database._DEFAULT = google.protobuf.symbol_database.SymbolDatabase(pool = google.protobuf.descriptor_pool.DescriptorPool()) - import caffe_pb2 as caffe_pb2 - google.protobuf.symbol_database._DEFAULT = old_symdb - sys.modules[__name__ + '.proto'] = sys.modules[__name__] - if shadow_caffe: - sys.modules['caffe'] = sys.modules[__name__] - sys.modules['caffe.proto'] = sys.modules[__name__] - return caffe_pb2 + +def initialize( + caffe_proto="https://raw.githubusercontent.com/BVLC/caffe/master/src/caffe/proto/caffe.proto", + codegen_dir=tempfile.mkdtemp(), + shadow_caffe=True, +): + global caffe_pb2 + if caffe_pb2 is None: + local_caffe_proto = os.path.join(codegen_dir, os.path.basename(caffe_proto)) + with open(local_caffe_proto, "w") as f: + mybytes = urlopen(caffe_proto).read() + mystr = mybytes.decode("ascii", "ignore") + f.write(mystr) + # f.write((urlopen if 'http' in caffe_proto else open)(caffe_proto).read()) + subprocess.check_call( + [ + "protoc", + "--proto_path", + os.path.dirname(local_caffe_proto), + "--python_out", + codegen_dir, + local_caffe_proto, + ] + ) + sys.path.insert(0, codegen_dir) + old_symdb = google.protobuf.symbol_database._DEFAULT + google.protobuf.symbol_database._DEFAULT = ( + google.protobuf.symbol_database.SymbolDatabase( + pool=google.protobuf.descriptor_pool.DescriptorPool() + ) + ) + import caffe_pb2 as caffe_pb2 + + google.protobuf.symbol_database._DEFAULT = old_symdb + sys.modules[__name__ + ".proto"] = sys.modules[__name__] + if shadow_caffe: + sys.modules["caffe"] = sys.modules[__name__] + sys.modules["caffe.proto"] = sys.modules[__name__] + return caffe_pb2 + def set_mode_gpu(): - global convert_to_gpu_if_enabled - convert_to_gpu_if_enabled = lambda obj: obj.cuda() + global convert_to_gpu_if_enabled + convert_to_gpu_if_enabled = lambda obj: obj.cuda() + def set_device(gpu_id): - torch.cuda.set_device(gpu_id) + torch.cuda.set_device(gpu_id) + class Net(nn.Module): - def __init__(self, prototxt, *args, **kwargs): - super(Net, self).__init__() - # to account for both constructors, see https://github.com/BVLC/caffe/blob/master/python/caffe/test/test_net.py#L145-L147 - caffe_proto = kwargs.pop('caffe_proto', None) - weights = kwargs.pop('weights', None) - phase = kwargs.pop('phase', None) - weights = weights or (args + (None, None))[0] - phase = phase or (args + (None, None))[1] - - self.net_param = initialize(caffe_proto).NetParameter() - google.protobuf.text_format.Parse(open(prototxt).read(), self.net_param) - - for layer in list(self.net_param.layer) + list(self.net_param.layers): - layer_type = layer.type if layer.type != 'Python' else layer.python_param.layer - if isinstance(layer_type, int): - layer_type = layer.LayerType.Name(layer_type) - module_constructor = ([v for k, v in modules.items() if k.replace('_', '').upper() in [layer_type.replace('_', '').upper(), layer.name.replace('_', '').upper()]] + [None])[0] - if module_constructor is not None: - param = to_dict(([v for f, v in layer.ListFields() if f.name.endswith('_param')] + [None])[0]) - caffe_input_variable_names = list(layer.bottom) - caffe_output_variable_names = list(layer.top) - caffe_loss_weight = (list(layer.loss_weight) or [1.0 if layer_type.upper().endswith('LOSS') else 0.0]) * len(layer.top) - caffe_propagate_down = list(getattr(layer, 'propagate_down', [])) or [True] * len(caffe_input_variable_names) - caffe_optimization_params = to_dict(layer.param) - param['inplace'] = len(caffe_input_variable_names) == 1 and caffe_input_variable_names == caffe_output_variable_names - module = module_constructor(param) - self.add_module(layer.name, module if isinstance(module, nn.Module) else CaffePythonLayerModule(module, caffe_input_variable_names, caffe_output_variable_names, param.get('param_str', '')) if type(module).__name__.endswith('Layer') else FunctionModule(module)) - module = getattr(self, layer.name) - module.caffe_layer_name = layer.name - module.caffe_layer_type = layer_type - module.caffe_input_variable_names = caffe_input_variable_names - module.caffe_output_variable_names = caffe_output_variable_names - module.caffe_loss_weight = caffe_loss_weight - module.caffe_propagate_down = caffe_propagate_down - module.caffe_optimization_params = caffe_optimization_params - for optim_param, p in zip(caffe_optimization_params, module.parameters()): - p.requires_grad = optim_param.get('lr_mult', 1) != 0 - else: - print('Skipping layer [{}, {}, {}]: not found in caffemodel2pytorch.modules dict'.format(layer.name, layer_type, layer.type)) - - if weights is not None: - self.copy_from(weights) - - self.blobs = collections.defaultdict(Blob) - self.blob_loss_weights = {name : loss_weight for module in self.children() for name, loss_weight in zip(module.caffe_output_variable_names, module.caffe_loss_weight)} - - self.train(phase != TEST) - convert_to_gpu_if_enabled(self) - - def forward(self, data = None, **variables): - if data is not None: - variables['data'] = data - numpy = not all(map(torch.is_tensor, variables.values())) - variables = {k : convert_to_gpu_if_enabled(torch.from_numpy(v.copy()) if numpy else v) for k, v in variables.items()} - - for module in [module for module in self.children() if not all(name in variables for name in module.caffe_output_variable_names)]: - for name in module.caffe_input_variable_names: - assert name in variables, 'Variable [{}] does not exist. Pass it as a keyword argument or provide a layer which produces it.'.format(name) - inputs = [variables[name] if propagate_down else variables[name].detach() for name, propagate_down in zip(module.caffe_input_variable_names, module.caffe_propagate_down)] - outputs = module(*inputs) - if not isinstance(outputs, tuple): - outputs = (outputs, ) - variables.update(dict(zip(module.caffe_output_variable_names, outputs))) - - self.blobs.update({k : Blob(data = v, numpy = numpy) for k, v in variables.items()}) - caffe_output_variable_names = set([name for module in self.children() for name in module.caffe_output_variable_names]) - set([name for module in self.children() for name in module.caffe_input_variable_names if name not in module.caffe_output_variable_names]) - return {k : v.detach().cpu().numpy() if numpy else v for k, v in variables.items() if k in caffe_output_variable_names} - - def copy_from(self, weights): - try: - import h5py, numpy - state_dict = self.state_dict() - for k, v in h5py.File(weights, 'r').items(): - if k in state_dict: - state_dict[k].resize_(v.shape).copy_(torch.from_numpy(numpy.array(v))) - print('caffemodel2pytorch: loaded model from [{}] in HDF5 format'.format(weights)) - except Exception as e: - print('caffemodel2pytorch: loading model from [{}] in HDF5 format failed [{}], falling back to caffemodel format'.format(weights, e)) - bytes_weights = open(weights, 'rb').read() - bytes_parsed = self.net_param.ParseFromString(bytes_weights) - if bytes_parsed != len(bytes_weights): - print('caffemodel2pytorch: loading model from [{}] in caffemodel format, WARNING: file length [{}] is not equal to number of parsed bytes [{}]'.format(weights, len(bytes_weights), bytes_parsed)) - for layer in list(self.net_param.layer) + list(self.net_param.layers): - module = getattr(self, layer.name, None) - if module is None: - continue - parameters = {name : convert_to_gpu_if_enabled(torch.FloatTensor(blob.data)).view(list(blob.shape.dim) if len(blob.shape.dim) > 0 else [blob.num, blob.channels, blob.height, blob.width]) for name, blob in zip(['weight', 'bias'], layer.blobs)} - if len(parameters) > 0: - module.set_parameters(**parameters) - print('caffemodel2pytorch: loaded model from [{}] in caffemodel format'.format(weights)) - - def save(self, weights): - import h5py - with h5py.File(weights, 'w') as h: - for k, v in self.state_dict().items(): - h[k] = v.cpu().numpy() - print('caffemodel2pytorch: saved model to [{}] in HDF5 format'.format(weights)) - - @property - def layers(self): - return list(self.children()) - -class Blob(object): - AssignmentAdapter = type('', (object, ), dict(shape = property(lambda self: self.contents.shape), __setitem__ = lambda self, indices, values: setattr(self, 'contents', values))) - - def __init__(self, data = None, diff = None, numpy = False): - self.data_ = data if data is not None else Blob.AssignmentAdapter() - self.diff_ = diff if diff is not None else Blob.AssignmentAdapter() - self.shape_ = None - self.numpy = numpy - - def reshape(self, *args): - self.shape_ = args - - def count(self, *axis): - return reduce(lambda x, y: x * y, self.shape_[slice(*(axis + [-1])[:2])]) - - @property - def data(self): - if self.numpy and torch.is_tensor(self.data_): - self.data_ = self.data_.detach().cpu().numpy() - return self.data_ - - @property - def diff(self): - if self.numpy and torch.is_tensor(self.diff_): - self.diff_ = self.diff_.detach().cpu().numpy() - return self.diff_ - - @property - def shape(self): - return self.shape_ if self.shape_ is not None else self.data_.shape - - @property - def num(self): - return self.shape[0] - - @property - def channels(self): - return self.shape[1] - - @property - def height(self): - return self.shape[2] - - @property - def width(self): - return self.shape[3] - -class Layer(torch.autograd.Function): - def __init__(self, caffe_python_layer = None, caffe_input_variable_names = None, caffe_output_variable_names = None, caffe_propagate_down = None): - self.caffe_python_layer = caffe_python_layer - self.caffe_input_variable_names = caffe_input_variable_names - self.caffe_output_variable_names = caffe_output_variable_names - self.caffe_propagate_down = caffe_propagate_down - - def forward(self, *inputs): - bottom = [Blob(data = v.cpu().numpy()) for v in inputs] - top = [Blob() for name in self.caffe_output_variable_names] - - #self.caffe_python_layer.reshape() - self.caffe_python_layer.setup(bottom, top) - self.caffe_python_layer.setup = lambda *args: None - - self.caffe_python_layer.forward(bottom, top) - outputs = tuple(convert_to_gpu_if_enabled(torch.from_numpy(v.data.contents.reshape(*v.shape))) for v in top) - self.save_for_backward(*(inputs + outputs)) - return outputs - - def backward(self, grad_outputs): - inputs, outputs = self.saved_tensors[:len(self.caffe_input_variable_names)], self.saved_tensors[len(self.caffe_input_variable_names):] - bottom = [Blob(data = v.cpu().numpy()) for v in inputs] - top = [Blob(data = output.cpu().numpy(), diff = grad_output.cpu().numpy()) for grad_output, output in zip(grad_outputs, outputs)] - self.caffe_python_layer.backward(top, self.caffe_propagate_down, bottom) - return tuple(convert_to_gpu_if_enabled(torch.from_numpy(blob.diff.contents.reshape(*v.reshape))) if propagate_down else None for v, propagate_down in zip(bottom, self.caffe_propagate_down)) - -class SGDSolver(object): - def __init__(self, solver_prototxt): - solver_param = initialize().SolverParameter() - google.protobuf.text_format.Parse(open(solver_prototxt).read(), solver_param) - solver_param = to_dict(solver_param) - self.net = Net(solver_param.get('train_net') or solver_param.get('net'), phase = TRAIN) - self.iter = 1 - self.iter_size = solver_param.get('iter_size', 1) - self.optimizer_params = dict(lr = solver_param.get('base_lr') / self.iter_size, momentum = solver_param.get('momentum', 0), weight_decay = solver_param.get('weight_decay', 0)) - self.lr_scheduler_params = dict(policy = solver_param.get('lr_policy'), step_size = solver_param.get('stepsize'), gamma = solver_param.get('gamma')) - self.optimizer, self.scheduler = None, None - - def init_optimizer_scheduler(self): - self.optimizer = torch.optim.SGD([dict(params = [param], lr = self.optimizer_params['lr'] * mult.get('lr_mult', 1), weight_decay = self.optimizer_params['weight_decay'] * mult.get('decay_mult', 1), momentum = self.optimizer_params['momentum']) for module in self.net.children() for param, mult in zip(module.parameters(), module.caffe_optimization_params + [{}, {}]) if param.requires_grad]) - self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size = self.lr_scheduler_params['step_size'], gamma = self.lr_scheduler_params['gamma']) if self.lr_scheduler_params.get('policy') == 'step' else type('', (object, ), dict(step = lambda self: None))() - - def step(self, iterations = 1, **inputs): - loss_total = 0.0 - for i in range(iterations): - tic = time.time() - if self.optimizer is not None: - self.optimizer.zero_grad() - - loss_batch = 0 - losses_batch = collections.defaultdict(float) - for j in range(self.iter_size): - outputs = [kv for kv in self.net(**inputs).items() if self.net.blob_loss_weights[kv[0]] != 0] - loss = sum([self.net.blob_loss_weights[k] * v.sum() for k, v in outputs]) - loss_batch += float(loss) / self.iter_size - for k, v in outputs: - losses_batch[k] += float(v.sum()) / self.iter_size - if self.optimizer is None: - self.init_optimizer_scheduler() - self.optimizer.zero_grad() - loss.backward() - - loss_total += loss_batch - self.optimizer.step() - self.scheduler.step() - self.iter += 1 - - log_prefix = self.__module__ + '.' + type(self).__name__ - print('{}] Iteration {}, loss: {}'.format(log_prefix, self.iter, loss_batch)) - for i, (name, loss) in enumerate(sorted(losses_batch.items())): - print('{}] Train net output #{}: {} = {} (* {} = {} loss)'.format(log_prefix, i, name, loss, self.net.blob_loss_weights[name], self.net.blob_loss_weights[name] * loss)) - print('{}] Iteration {}, lr = {}, time = {}'.format(log_prefix, self.iter, self.optimizer_params['lr'], time.time() - tic)) - - return loss_total - -modules = dict( - Convolution = lambda param: Convolution(param), - InnerProduct = lambda param: InnerProduct(param), - Pooling = lambda param: [nn.MaxPool2d, nn.AvgPool2d][param['pool']](kernel_size = first_or(param, 'kernel_size', 1), stride = first_or(param, 'stride', 1), padding = first_or(param, 'pad', 0)), - Softmax = lambda param: nn.Softmax(dim = param.get('axis', -1)), - ReLU = lambda param: nn.ReLU(), - Dropout = lambda param: nn.Dropout(p = param['dropout_ratio']), - Eltwise = lambda param: [torch.mul, torch.add, torch.max][param.get('operation', 1)], - LRN = lambda param: nn.LocalResponseNorm(size = param['local_size'], alpha = param['alpha'], beta = param['beta']) -) + def __init__(self, prototxt, *args, **kwargs): + super().__init__() + # to account for both constructors, see https://github.com/BVLC/caffe/blob/master/python/caffe/test/test_net.py#L145-L147 + caffe_proto = kwargs.pop("caffe_proto", None) + weights = kwargs.pop("weights", None) + print(weights) + phase = kwargs.pop("phase", None) + weights = weights or (args + (None, None))[0] + phase = phase or (args + (None, None))[1] + + self.net_param = initialize(caffe_proto).NetParameter() + google.protobuf.text_format.Parse(open(prototxt).read(), self.net_param) + + # Check if number_of_inputs is not larger than 1 + number_of_inputs = len(self.net_param.input_shape) + if number_of_inputs > 1: + raise NotImplementedError("Multiple inputs is currently not supported.") + + # Check dimensionality + # The first two dimensions are batch and channel, so subtract 2 + input_dimensions = len(self.net_param.input_shape[0].dim) - 2 + if not input_dimensions in [2, 3]: + raise NotImplementedError( + "Only 2D and 3D caffe networks are currently supported." + ) + + dimensions = 2 + if dimensions == 2: + modules = modules2d + + for layer in list(self.net_param.layer) + list(self.net_param.layers): + layer_type = ( + layer.type if layer.type != "Python" else layer.python_param.layer + ) + + if isinstance(layer_type, int): + layer_type = layer.LayerType.Name(layer_type) + + module_constructor = ( + [ + v + for k, v in modules.items() + if k.replace("_", "").upper() + in [ + layer_type.replace("_", "").upper(), + layer.name.replace("_", "").upper(), + ] + ] + + [None] + )[0] + + print(f"Module constructor: {module_constructor}") + + if module_constructor is not None: + param = to_dict( + ( + [v for f, v in layer.ListFields() if f.name.endswith("_param")] + + [None] + )[0] + ) + caffe_input_variable_names = list(layer.bottom) + caffe_output_variable_names = list(layer.top) + caffe_loss_weight = ( + list(layer.loss_weight) + or [1.0 if layer_type.upper().endswith("LOSS") else 0.0] + ) * len(layer.top) + caffe_propagate_down = list(getattr(layer, "propagate_down", [])) or [ + True + ] * len(caffe_input_variable_names) + caffe_optimization_params = to_dict(layer.param) + param["inplace"] = ( + len(caffe_input_variable_names) == 1 + and caffe_input_variable_names == caffe_output_variable_names + ) + module = module_constructor(param) + self.add_module( + layer.name, + module + if isinstance(module, nn.Module) + else CaffePythonLayerModule( + module, + caffe_input_variable_names, + caffe_output_variable_names, + param.get("param_str", ""), + ) + if type(module).__name__.endswith("Layer") + else FunctionModule(module), + ) + module = getattr(self, layer.name) + module.caffe_layer_name = layer.name + module.caffe_layer_type = layer_type + module.caffe_input_variable_names = caffe_input_variable_names + module.caffe_output_variable_names = caffe_output_variable_names + module.caffe_loss_weight = caffe_loss_weight + module.caffe_propagate_down = caffe_propagate_down + module.caffe_optimization_params = caffe_optimization_params + for optim_param, param in zip( + caffe_optimization_params, module.parameters() + ): + param.requires_grad = optim_param.get("lr_mult", 1) != 0 + else: + print( + ( + f"Skipping layer [{layer.name}, {layer_type}, {layer.type}]:" + "not found in caffemodel2pytorch.modules dict" + ) + ) + + if weights is not None: + self.copy_from(weights) + + self.blobs = collections.defaultdict(Blob) + self.blob_loss_weights = { + name: loss_weight + for module in self.children() + for name, loss_weight in zip( + module.caffe_output_variable_names, module.caffe_loss_weight + ) + } + + self.train(phase != TEST) + convert_to_gpu_if_enabled(self) + + def forward(self, data=None, **variables): + if data is not None: + variables["data"] = data + numpy = not all(map(torch.is_tensor, variables.values())) + variables = { + k: convert_to_gpu_if_enabled(torch.from_numpy(v.copy()) if numpy else v) + for k, v in variables.items() + } + + for module in [ + module + for module in self.children() + if not all(name in variables for name in module.caffe_output_variable_names) + ]: + for name in module.caffe_input_variable_names: + assert name in variables, ( + f"Variable [{name}] does not exist. " + "Pass it as a keyword argument or provide a layer which produces it." + ) + inputs = [ + variables[name] if propagate_down else variables[name].detach() + for name, propagate_down in zip( + module.caffe_input_variable_names, module.caffe_propagate_down + ) + ] + outputs = module(*inputs) + if not isinstance(outputs, tuple): + outputs = (outputs,) + variables.update(dict(zip(module.caffe_output_variable_names, outputs))) + + self.blobs.update({k: Blob(data=v, numpy=numpy) for k, v in variables.items()}) + caffe_output_variable_names = set( + [ + name + for module in self.children() + for name in module.caffe_output_variable_names + ] + ) - set( + [ + name + for module in self.children() + for name in module.caffe_input_variable_names + if name not in module.caffe_output_variable_names + ] + ) + return { + k: v.detach().cpu().numpy() if numpy else v + for k, v in variables.items() + if k in caffe_output_variable_names + } + + def copy_from(self, weights): + try: + import h5py, numpy + + state_dict = self.state_dict() + for k, v in h5py.File(weights, "r").items(): + if k in state_dict: + state_dict[k].resize_(v.shape).copy_( + torch.from_numpy(numpy.array(v)) + ) + print( + "caffemodel2pytorch: loaded model from [{}] in HDF5 format".format( + weights + ) + ) + except Exception as e: + print( + "caffemodel2pytorch: loading model from [{}] in HDF5 format failed [{}], falling back to caffemodel format".format( + weights, e + ) + ) + bytes_weights = open(weights, "rb").read() + bytes_parsed = self.net_param.ParseFromString(bytes_weights) + if bytes_parsed != len(bytes_weights): + print( + "caffemodel2pytorch: loading model from [{}] in caffemodel format, WARNING: file length [{}] is not equal to number of parsed bytes [{}]".format( + weights, len(bytes_weights), bytes_parsed + ) + ) + for layer in list(self.net_param.layer) + list(self.net_param.layers): + module = getattr(self, layer.name, None) + if module is None: + continue + parameters = { + name: convert_to_gpu_if_enabled(torch.FloatTensor(blob.data)).view( + list(blob.shape.dim) + if len(blob.shape.dim) > 0 + else [blob.num, blob.channels, blob.height, blob.width] + ) + for name, blob in zip(["weight", "bias"], layer.blobs) + } + if len(parameters) > 0: + print(f"Weights shape: {parameters['weight'].shape}") + print(f"Bias: {parameters['bias']}") + module.set_parameters(**parameters) + print( + "caffemodel2pytorch: loaded model from [{}] in caffemodel format".format( + weights + ) + ) + + def save(self, weights): + import h5py + + with h5py.File(weights, "w") as h: + for k, v in self.state_dict().items(): + h[k] = v.cpu().numpy() + print("caffemodel2pytorch: saved model to [{}] in HDF5 format".format(weights)) + + @property + def layers(self): + return list(self.children()) + class FunctionModule(nn.Module): - def __init__(self, forward): - super(FunctionModule, self).__init__() - self.forward_func = forward + def __init__(self, forward): + super(FunctionModule, self).__init__() + self.forward_func = forward + + def forward(self, *inputs): + return self.forward_func(*inputs) - def forward(self, *inputs): - return self.forward_func(*inputs) class CaffePythonLayerModule(nn.Module): - def __init__(self, caffe_python_layer, caffe_input_variable_names, caffe_output_variable_names, param_str): - super(CaffePythonLayerModule, self).__init__() - caffe_python_layer.param_str = param_str - self.caffe_python_layer = caffe_python_layer - self.caffe_input_variable_names = caffe_input_variable_names - self.caffe_output_variable_names = caffe_output_variable_names - - def forward(self, *inputs): - return Layer(self.caffe_python_layer, self.caffe_input_variable_names, self.caffe_output_variable_names)(*inputs) - - def __getattr__(self, name): - return nn.Module.__getattr__(self, name) if name in dir(self) else getattr(self.caffe_python_layer, name) - -class Convolution(nn.Conv2d): - def __init__(self, param): - super(Convolution, self).__init__(first_or(param,'group',1), param['num_output'], kernel_size = first_or(param, 'kernel_size', 1), stride = first_or(param, 'stride', 1), padding = first_or(param, 'pad', 0), dilation = first_or(param, 'dilation', 1), groups = first_or(param, 'group', 1)) - self.weight, self.bias = nn.Parameter(), nn.Parameter() - self.weight_init, self.bias_init = param.get('weight_filler', {}), param.get('bias_filler', {}) - - def forward(self, x): - if self.weight.numel() == 0 and self.bias.numel() == 0: - requires_grad = [self.weight.requires_grad, self.bias.requires_grad] - super(Convolution, self).__init__(x.size(1), self.out_channels, kernel_size = self.kernel_size, stride = self.stride, padding = self.padding, dilation = self.dilation) - convert_to_gpu_if_enabled(self) - init_weight_bias(self, requires_grad = requires_grad) - return super(Convolution, self).forward(x) - - def set_parameters(self, weight = None, bias = None): - init_weight_bias(self, weight = weight, bias = bias.view(-1) if bias is not None else bias) - self.in_channels = self.weight.size(1) - -class InnerProduct(nn.Linear): - def __init__(self, param): - super(InnerProduct, self).__init__(1, param['num_output']) - self.weight, self.bias = nn.Parameter(), nn.Parameter() - self.weight_init, self.bias_init = param.get('weight_filler', {}), param.get('bias_filler', {}) - - def forward(self, x): - if self.weight.numel() == 0 and self.bias.numel() == 0: - requires_grad = [self.weight.requires_grad, self.bias.requires_grad] - super(InnerProduct, self).__init__(x.size(1), self.out_features) - convert_to_gpu_if_enabled(self) - init_weight_bias(self, requires_grad = requires_grad) - return super(InnerProduct, self).forward(x if x.size(-1) == self.in_features else x.view(len(x), -1)) - - def set_parameters(self, weight = None, bias = None): - init_weight_bias(self, weight = weight.view(weight.size(-2), weight.size(-1)) if weight is not None else None, bias = bias.view(-1) if bias is not None else None) - self.in_features = self.weight.size(1) - -def init_weight_bias(self, weight = None, bias = None, requires_grad = []): - if weight is not None: - self.weight = nn.Parameter(weight.type_as(self.weight), requires_grad = self.weight.requires_grad) - if bias is not None: - self.bias = nn.Parameter(bias.type_as(self.bias), requires_grad = self.bias.requires_grad) - for name, requires_grad in zip(['weight', 'bias'], requires_grad): - param, init = getattr(self, name), getattr(self, name + '_init') - if init.get('type') == 'gaussian': - nn.init.normal_(param, std = init['std']) - elif init.get('type') == 'constant': - nn.init.constant_(param, val = init['value']) - param.requires_grad = requires_grad - -def convert_to_gpu_if_enabled(obj): - return obj - -def first_or(param, key, default): - return param[key] if isinstance(param.get(key), int) else (param.get(key, []) + [default])[0] - -def to_dict(obj): - return list(map(to_dict, obj)) if isinstance(obj, collections.Iterable) else {} if obj is None else {f.name : converter(v) if f.label != FD.LABEL_REPEATED else list(map(converter, v)) for f, v in obj.ListFields() for converter in [{FD.TYPE_DOUBLE: float, FD.TYPE_SFIXED32: float, FD.TYPE_SFIXED64: float, FD.TYPE_SINT32: int, FD.TYPE_SINT64: int, FD.TYPE_FLOAT: float, FD.TYPE_ENUM: int, FD.TYPE_UINT32: int, FD.TYPE_INT64: int, FD.TYPE_UINT64: int, FD.TYPE_INT32: int, FD.TYPE_FIXED64: float, FD.TYPE_FIXED32: float, FD.TYPE_BOOL: bool, FD.TYPE_STRING: str, FD.TYPE_BYTES: lambda x: x.encode('string_escape'), FD.TYPE_MESSAGE: to_dict}[f.type]]} - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument(metavar = 'model.caffemodel', dest = 'model_caffemodel', help = 'Path to model.caffemodel') - parser.add_argument('-o', dest = 'output_path', help = 'Path to converted model, supported file extensions are: h5, npy, npz, json, pt') - parser.add_argument('--caffe.proto', metavar = '--caffe.proto', dest = 'caffe_proto', help = 'Path to caffe.proto (typically located at CAFFE_ROOT/src/caffe/proto/caffe.proto)', default = 'https://raw.githubusercontent.com/BVLC/caffe/master/src/caffe/proto/caffe.proto') - args = parser.parse_args() - args.output_path = args.output_path or args.model_caffemodel + '.pt' - - net_param = initialize(args.caffe_proto).NetParameter() - net_param.ParseFromString(open(args.model_caffemodel, 'rb').read()) - blobs = {layer.name + '.' + name : dict(data = blob.data, shape = list(blob.shape.dim) if len(blob.shape.dim) > 0 else [blob.num, blob.channels, blob.height, blob.width]) for layer in list(net_param.layer) + list(net_param.layers) for name, blob in zip(['weight', 'bias'], layer.blobs)} - - if args.output_path.endswith('.json'): - import json - with open(args.output_path, 'w') as f: - json.dump(blobs, f) - elif args.output_path.endswith('.h5'): - import h5py, numpy - with h5py.File(args.output_path, 'w') as h: - h.update(**{k : numpy.array(blob['data'], dtype = numpy.float32).reshape(*blob['shape']) for k, blob in blobs.items()}) - elif args.output_path.endswith('.npy') or args.output_path.endswith('.npz'): - import numpy - (numpy.savez if args.output_path[-1] == 'z' else numpy.save)(args.output_path, **{k : numpy.array(blob['data'], dtype = numpy.float32).reshape(*blob['shape']) for k, blob in blobs.items()}) - elif args.output_path.endswith('.pt'): - torch.save({k : torch.FloatTensor(blob['data']).view(*blob['shape']) for k, blob in blobs.items()}, args.output_path) + def __init__( + self, + caffe_python_layer, + caffe_input_variable_names, + caffe_output_variable_names, + param_str, + ): + super(CaffePythonLayerModule, self).__init__() + caffe_python_layer.param_str = param_str + self.caffe_python_layer = caffe_python_layer + self.caffe_input_variable_names = caffe_input_variable_names + self.caffe_output_variable_names = caffe_output_variable_names + + def forward(self, *inputs): + return Layer( + self.caffe_python_layer, + self.caffe_input_variable_names, + self.caffe_output_variable_names, + )(*inputs) + + def __getattr__(self, name): + return ( + nn.Module.__getattr__(self, name) + if name in dir(self) + else getattr(self.caffe_python_layer, name) + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + metavar="model.caffemodel", + dest="model_caffemodel", + help="Path to model.caffemodel", + ) + parser.add_argument( + "-o", + dest="output_path", + help="Path to converted model, supported file extensions are: h5, npy, npz, json, pt", + ) + parser.add_argument( + "--caffe.proto", + metavar="--caffe.proto", + dest="caffe_proto", + help="Path to caffe.proto (typically located at CAFFE_ROOT/src/caffe/proto/caffe.proto)", + default="https://raw.githubusercontent.com/BVLC/caffe/master/src/caffe/proto/caffe.proto", + ) + args = parser.parse_args() + args.output_path = args.output_path or args.model_caffemodel + ".pt" + + net_param = initialize(args.caffe_proto).NetParameter() + net_param.ParseFromString(open(args.model_caffemodel, "rb").read()) + blobs = { + layer.name + + "." + + name: dict( + data=blob.data, + shape=list(blob.shape.dim) + if len(blob.shape.dim) > 0 + else [blob.num, blob.channels, blob.height, blob.width], + ) + for layer in list(net_param.layer) + list(net_param.layers) + for name, blob in zip(["weight", "bias"], layer.blobs) + } + + if args.output_path.endswith(".json"): + import json + + with open(args.output_path, "w") as f: + json.dump(blobs, f) + elif args.output_path.endswith(".h5"): + import h5py, numpy + + with h5py.File(args.output_path, "w") as h: + h.update( + **{ + k: numpy.array(blob["data"], dtype=numpy.float32).reshape( + *blob["shape"] + ) + for k, blob in blobs.items() + } + ) + elif args.output_path.endswith(".npy") or args.output_path.endswith(".npz"): + import numpy + + (numpy.savez if args.output_path[-1] == "z" else numpy.save)( + args.output_path, + **{ + k: numpy.array(blob["data"], dtype=numpy.float32).reshape( + *blob["shape"] + ) + for k, blob in blobs.items() + }, + ) + elif args.output_path.endswith(".pt"): + torch.save( + { + k: torch.FloatTensor(blob["data"]).view(*blob["shape"]) + for k, blob in blobs.items() + }, + args.output_path, + ) diff --git a/layers2d.py b/layers2d.py new file mode 100644 index 0000000..ea1c3e4 --- /dev/null +++ b/layers2d.py @@ -0,0 +1,291 @@ +import torch +from torch import nn + +from blob import Blob +from utils import first_or, init_weight_bias + + +class Layer(torch.autograd.Function): + def __init__( + self, + caffe_python_layer=None, + caffe_input_variable_names=None, + caffe_output_variable_names=None, + caffe_propagate_down=None, + ): + self.caffe_python_layer = caffe_python_layer + self.caffe_input_variable_names = caffe_input_variable_names + self.caffe_output_variable_names = caffe_output_variable_names + self.caffe_propagate_down = caffe_propagate_down + + def forward(self, *inputs): + bottom = [Blob(data=v.cpu().numpy()) for v in inputs] + top = [Blob() for name in self.caffe_output_variable_names] + + # self.caffe_python_layer.reshape() + self.caffe_python_layer.setup(bottom, top) + self.caffe_python_layer.setup = lambda *args: None + + self.caffe_python_layer.forward(bottom, top) + outputs = tuple( + convert_to_gpu_if_enabled( + torch.from_numpy(v.data.contents.reshape(*v.shape)) + ) + for v in top + ) + self.save_for_backward(*(inputs + outputs)) + return outputs + + def backward(self, grad_outputs): + inputs, outputs = ( + self.saved_tensors[: len(self.caffe_input_variable_names)], + self.saved_tensors[len(self.caffe_input_variable_names) :], + ) + bottom = [Blob(data=v.cpu().numpy()) for v in inputs] + top = [ + Blob(data=output.cpu().numpy(), diff=grad_output.cpu().numpy()) + for grad_output, output in zip(grad_outputs, outputs) + ] + self.caffe_python_layer.backward(top, self.caffe_propagate_down, bottom) + return tuple( + convert_to_gpu_if_enabled( + torch.from_numpy(blob.diff.contents.reshape(*v.reshape)) + ) + if propagate_down + else None + for v, propagate_down in zip(bottom, self.caffe_propagate_down) + ) + + +class SGDSolver(object): + def __init__(self, solver_prototxt): + solver_param = initialize().SolverParameter() + google.protobuf.text_format.Parse(open(solver_prototxt).read(), solver_param) + solver_param = to_dict(solver_param) + self.net = Net( + solver_param.get("train_net") or solver_param.get("net"), phase=TRAIN + ) + self.iter = 1 + self.iter_size = solver_param.get("iter_size", 1) + self.optimizer_params = dict( + lr=solver_param.get("base_lr") / self.iter_size, + momentum=solver_param.get("momentum", 0), + weight_decay=solver_param.get("weight_decay", 0), + ) + self.lr_scheduler_params = dict( + policy=solver_param.get("lr_policy"), + step_size=solver_param.get("stepsize"), + gamma=solver_param.get("gamma"), + ) + self.optimizer, self.scheduler = None, None + + def init_optimizer_scheduler(self): + self.optimizer = torch.optim.SGD( + [ + dict( + params=[param], + lr=self.optimizer_params["lr"] * mult.get("lr_mult", 1), + weight_decay=self.optimizer_params["weight_decay"] + * mult.get("decay_mult", 1), + momentum=self.optimizer_params["momentum"], + ) + for module in self.net.children() + for param, mult in zip( + module.parameters(), module.caffe_optimization_params + [{}, {}] + ) + if param.requires_grad + ] + ) + self.scheduler = ( + torch.optim.lr_scheduler.StepLR( + self.optimizer, + step_size=self.lr_scheduler_params["step_size"], + gamma=self.lr_scheduler_params["gamma"], + ) + if self.lr_scheduler_params.get("policy") == "step" + else type("", (object,), dict(step=lambda self: None))() + ) + + def step(self, iterations=1, **inputs): + loss_total = 0.0 + for i in range(iterations): + tic = time.time() + if self.optimizer is not None: + self.optimizer.zero_grad() + + loss_batch = 0 + losses_batch = collections.defaultdict(float) + for j in range(self.iter_size): + outputs = [ + kv + for kv in self.net(**inputs).items() + if self.net.blob_loss_weights[kv[0]] != 0 + ] + loss = sum( + [self.net.blob_loss_weights[k] * v.sum() for k, v in outputs] + ) + loss_batch += float(loss) / self.iter_size + for k, v in outputs: + losses_batch[k] += float(v.sum()) / self.iter_size + if self.optimizer is None: + self.init_optimizer_scheduler() + self.optimizer.zero_grad() + loss.backward() + + loss_total += loss_batch + self.optimizer.step() + self.scheduler.step() + self.iter += 1 + + log_prefix = self.__module__ + "." + type(self).__name__ + print( + "{}] Iteration {}, loss: {}".format(log_prefix, self.iter, loss_batch) + ) + for i, (name, loss) in enumerate(sorted(losses_batch.items())): + print( + "{}] Train net output #{}: {} = {} (* {} = {} loss)".format( + log_prefix, + i, + name, + loss, + self.net.blob_loss_weights[name], + self.net.blob_loss_weights[name] * loss, + ) + ) + print( + "{}] Iteration {}, lr = {}, time = {}".format( + log_prefix, + self.iter, + self.optimizer_params["lr"], + time.time() - tic, + ) + ) + + return loss_total + + +class Convolution(nn.Conv2d): + def __init__(self, param): + super(Convolution, self).__init__( + first_or(param, "group", 1), + param["num_output"], + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + dilation=first_or(param, "dilation", 1), + groups=first_or(param, "group", 1), + ) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight.numel() == 0 and self.bias.numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(Convolution, self).__init__( + x.size(1), + self.out_channels, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + ) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(Convolution, self).forward(x) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, weight=weight, bias=bias.view(-1) if bias is not None else bias + ) + self.in_channels = self.weight.size(1) + + +class Deconvolution(nn.ConvTranspose2d): + def __init__(self, param): + super(Deconvolution, self).__init__( + first_or(param, "group", 1), + param["num_output"], + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + dilation=first_or(param, "dilation", 1), + groups=first_or(param, "group", 1), + ) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight_numel() == 0 and self.bias_numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(Deconvolution, self).__init__( + x.size(1), + self.out_channels, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + ) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(Deconvolution, self).forward(x) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, weight=weight, bias=bias.view(-1) if bias is not None else bias + ) + self.in_channels = self.weight.size(1) + + +class InnerProduct(nn.Linear): + def __init__(self, param): + super(InnerProduct, self).__init__(1, param["num_output"]) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight.numel() == 0 and self.bias.numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(InnerProduct, self).__init__(x.size(1), self.out_features) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(InnerProduct, self).forward( + x if x.size(-1) == self.in_features else x.view(len(x), -1) + ) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, + weight=weight.view(weight.size(-2), weight.size(-1)) + if weight is not None + else None, + bias=bias.view(-1) if bias is not None else None, + ) + self.in_features = self.weight.size(1) + + +# using dict calls (=) otherwise the dict does not play nice with +# the lambda function +modules2d = dict( + Convolution=lambda param: Convolution(param), + Deconvolution=lambda param: Deconvolution(param), + InnerProduct=lambda param: InnerProduct(param), + Pooling=lambda param: [nn.MaxPool2d, nn.AvgPool2d][param["pool"]]( + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + ), + Softmax=lambda param: nn.Softmax(dim=param.get("axis", -1)), + ReLU=lambda param: nn.ReLU(), + Dropout=lambda param: nn.Dropout(p=param["dropout_ratio"]), + Eltwise=lambda param: [torch.mul, torch.add, torch.max][param.get("operation", 1)], + Concat=lambda param: torch.cat, + LRN=lambda param: nn.LocalResponseNorm( + size=param["local_size"], alpha=param["alpha"], beta=param["beta"] + ), +) diff --git a/test.py b/test.py new file mode 100644 index 0000000..994b6d8 --- /dev/null +++ b/test.py @@ -0,0 +1,41 @@ +import torch +import caffemodel2pytorch + +from torchsummary import summary +import skimage + +import SimpleITK as sitk + +import matplotlib.pyplot as plt + +model = caffemodel2pytorch.Net( + prototxt="../CDSR/x2/CDSRx2.prototxt", + weights="../CDSR/x2/solver1_iter_104192.caffemodel", + caffe_proto="http://raw.githubusercontent.com/BVLC/caffe/master/src/caffe/proto/caffe.proto", +) +print(model) + +img = sitk.ReadImage("../CDSR/test_data/KKI2009-01-MPRAGE.nii") +img_arr = sitk.GetArrayFromImage(img) + +img_arr -= img_arr.min() +img_arr /= img_arr.max() + +img_down2 = skimage.transform.rescale(img_arr, scale=0.5, anti_aliasing=True) +img_down2 -= img_down2.min() +img_down2 /= img_down2.max() +img_up2 = skimage.transform.rescale(img_down2, scale=2.0, anti_aliasing=True, order=1) + +img_down3 = skimage.transform.rescale(img_arr, scale=0.3333333, anti_aliasing=True) +img_down3 -= img_down3.min() +img_down3 /= img_down3.max() +img_up3 = skimage.transform.rescale(img_down3, scale=3.0, anti_aliasing=True, order=1) + +super_res = model(img_up2[None, None, ...]) + +fig, ax = plt.subplots(nrows=2, ncols=2) +ax[0, 0].imshow(img_arr[170, ...]) +ax[0, 1].imshow(img_up2[170, ...]) +ax[1, 0].imshow(img_up3[170, ...]) +ax[1, 1].imshow(super_res) +plt.show() diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..bfb7e4b --- /dev/null +++ b/utils.py @@ -0,0 +1,72 @@ +from torch import nn +import collections + +from google.protobuf.descriptor import FieldDescriptor as FD + + +def init_weight_bias(self, weight=None, bias=None, requires_grad=[]): + if weight is not None: + self.weight = nn.Parameter( + weight.type_as(self.weight), requires_grad=self.weight.requires_grad + ) + if bias is not None: + self.bias = nn.Parameter( + bias.type_as(self.bias), requires_grad=self.bias.requires_grad + ) + for name, requires_grad in zip(["weight", "bias"], requires_grad): + param, init = getattr(self, name), getattr(self, name + "_init") + if init.get("type") == "gaussian": + nn.init.normal_(param, std=init["std"]) + elif init.get("type") == "constant": + nn.init.constant_(param, val=init["value"]) + param.requires_grad = requires_grad + + +def convert_to_gpu_if_enabled(obj): + return obj + + +def first_or(param, key, default): + val = param.get(key, None) + print(f"{key}, {val}, {default}") + if not val: + return default + if isinstance(val, float): + return val + return val + + +def to_dict(obj): + return ( + list(map(to_dict, obj)) + if isinstance(obj, collections.Iterable) + else {} + if obj is None + else { + f.name: converter(v) + if f.label != FD.LABEL_REPEATED + else list(map(converter, v)) + for f, v in obj.ListFields() + for converter in [ + { + FD.TYPE_DOUBLE: float, + FD.TYPE_SFIXED32: float, + FD.TYPE_SFIXED64: float, + FD.TYPE_SINT32: int, + FD.TYPE_SINT64: int, + FD.TYPE_FLOAT: float, + FD.TYPE_ENUM: int, + FD.TYPE_UINT32: int, + FD.TYPE_INT64: int, + FD.TYPE_UINT64: int, + FD.TYPE_INT32: int, + FD.TYPE_FIXED64: float, + FD.TYPE_FIXED32: float, + FD.TYPE_BOOL: bool, + FD.TYPE_STRING: str, + FD.TYPE_BYTES: lambda x: x.encode("string_escape"), + FD.TYPE_MESSAGE: to_dict, + }[f.type] + ] + } + ) From b212eeed9550d59b4a11fddf28093edb6c74095c Mon Sep 17 00:00:00 2001 From: dirk Date: Wed, 29 Mar 2023 15:02:11 +0200 Subject: [PATCH 4/7] Add 3D support --- caffemodel2pytorch.py | 9 +- layers3d.py | 296 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 303 insertions(+), 2 deletions(-) create mode 100644 layers3d.py diff --git a/caffemodel2pytorch.py b/caffemodel2pytorch.py index 4a8e2e2..2f765ed 100755 --- a/caffemodel2pytorch.py +++ b/caffemodel2pytorch.py @@ -20,6 +20,7 @@ import ssl from layers2d import modules2d +from layers3d import modules3d from blob import Blob from utils import to_dict, convert_to_gpu_if_enabled @@ -108,9 +109,13 @@ def __init__(self, prototxt, *args, **kwargs): "Only 2D and 3D caffe networks are currently supported." ) - dimensions = 2 - if dimensions == 2: + # Decide to use 2D or 3D modules + if input_dimensions == 2: + print("2D modules selected") modules = modules2d + if input_dimensions == 3: + print("3D modules selected") + modules = modules3d for layer in list(self.net_param.layer) + list(self.net_param.layers): layer_type = ( diff --git a/layers3d.py b/layers3d.py new file mode 100644 index 0000000..855ad8b --- /dev/null +++ b/layers3d.py @@ -0,0 +1,296 @@ +import torch +from torch import nn + +from blob import Blob +from utils import first_or, init_weight_bias, convert_to_gpu_if_enabled + + +class Layer(torch.autograd.Function): + def __init__( + self, + caffe_python_layer=None, + caffe_input_variable_names=None, + caffe_output_variable_names=None, + caffe_propagate_down=None, + ): + self.caffe_python_layer = caffe_python_layer + self.caffe_input_variable_names = caffe_input_variable_names + self.caffe_output_variable_names = caffe_output_variable_names + self.caffe_propagate_down = caffe_propagate_down + + def forward(self, *inputs): + bottom = [Blob(data=v.cpu().numpy()) for v in inputs] + top = [Blob() for name in self.caffe_output_variable_names] + + # self.caffe_python_layer.reshape() + self.caffe_python_layer.setup(bottom, top) + self.caffe_python_layer.setup = lambda *args: None + + self.caffe_python_layer.forward(bottom, top) + outputs = tuple( + convert_to_gpu_if_enabled( + torch.from_numpy(v.data.contents.reshape(*v.shape)) + ) + for v in top + ) + self.save_for_backward(*(inputs + outputs)) + return outputs + + def backward(self, grad_outputs): + inputs, outputs = ( + self.saved_tensors[: len(self.caffe_input_variable_names)], + self.saved_tensors[len(self.caffe_input_variable_names) :], + ) + bottom = [Blob(data=v.cpu().numpy()) for v in inputs] + top = [ + Blob(data=output.cpu().numpy(), diff=grad_output.cpu().numpy()) + for grad_output, output in zip(grad_outputs, outputs) + ] + self.caffe_python_layer.backward(top, self.caffe_propagate_down, bottom) + return tuple( + convert_to_gpu_if_enabled( + torch.from_numpy(blob.diff.contents.reshape(*v.reshape)) + ) + if propagate_down + else None + for v, propagate_down in zip(bottom, self.caffe_propagate_down) + ) + + +class SGDSolver(object): + def __init__(self, solver_prototxt): + solver_param = initialize().SolverParameter() + google.protobuf.text_format.Parse(open(solver_prototxt).read(), solver_param) + solver_param = to_dict(solver_param) + self.net = Net( + solver_param.get("train_net") or solver_param.get("net"), phase=TRAIN + ) + self.iter = 1 + self.iter_size = solver_param.get("iter_size", 1) + self.optimizer_params = dict( + lr=solver_param.get("base_lr") / self.iter_size, + momentum=solver_param.get("momentum", 0), + weight_decay=solver_param.get("weight_decay", 0), + ) + self.lr_scheduler_params = dict( + policy=solver_param.get("lr_policy"), + step_size=solver_param.get("stepsize"), + gamma=solver_param.get("gamma"), + ) + self.optimizer, self.scheduler = None, None + + def init_optimizer_scheduler(self): + self.optimizer = torch.optim.SGD( + [ + dict( + params=[param], + lr=self.optimizer_params["lr"] * mult.get("lr_mult", 1), + weight_decay=self.optimizer_params["weight_decay"] + * mult.get("decay_mult", 1), + momentum=self.optimizer_params["momentum"], + ) + for module in self.net.children() + for param, mult in zip( + module.parameters(), module.caffe_optimization_params + [{}, {}] + ) + if param.requires_grad + ] + ) + self.scheduler = ( + torch.optim.lr_scheduler.StepLR( + self.optimizer, + step_size=self.lr_scheduler_params["step_size"], + gamma=self.lr_scheduler_params["gamma"], + ) + if self.lr_scheduler_params.get("policy") == "step" + else type("", (object,), dict(step=lambda self: None))() + ) + + def step(self, iterations=1, **inputs): + loss_total = 0.0 + for i in range(iterations): + tic = time.time() + if self.optimizer is not None: + self.optimizer.zero_grad() + + loss_batch = 0 + losses_batch = collections.defaultdict(float) + for j in range(self.iter_size): + outputs = [ + kv + for kv in self.net(**inputs).items() + if self.net.blob_loss_weights[kv[0]] != 0 + ] + loss = sum( + [self.net.blob_loss_weights[k] * v.sum() for k, v in outputs] + ) + loss_batch += float(loss) / self.iter_size + for k, v in outputs: + losses_batch[k] += float(v.sum()) / self.iter_size + if self.optimizer is None: + self.init_optimizer_scheduler() + self.optimizer.zero_grad() + loss.backward() + + loss_total += loss_batch + self.optimizer.step() + self.scheduler.step() + self.iter += 1 + + log_prefix = self.__module__ + "." + type(self).__name__ + print( + "{}] Iteration {}, loss: {}".format(log_prefix, self.iter, loss_batch) + ) + for i, (name, loss) in enumerate(sorted(losses_batch.items())): + print( + "{}] Train net output #{}: {} = {} (* {} = {} loss)".format( + log_prefix, + i, + name, + loss, + self.net.blob_loss_weights[name], + self.net.blob_loss_weights[name] * loss, + ) + ) + print( + "{}] Iteration {}, lr = {}, time = {}".format( + log_prefix, + self.iter, + self.optimizer_params["lr"], + time.time() - tic, + ) + ) + + return loss_total + + +class Convolution(nn.Conv3d): + def __init__(self, param): + super(Convolution, self).__init__( + first_or(param, "group", 1), + param["num_output"], + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + dilation=first_or(param, "dilation", 1), + groups=first_or(param, "group", 1), + ) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight.numel() == 0 and self.bias.numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(Convolution, self).__init__( + x.size(1), + self.out_channels, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + ) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(Convolution, self).forward(x) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, weight=weight, bias=bias.view(-1) if bias is not None else bias + ) + self.in_channels = self.weight.size(1) + + +class Deconvolution(nn.ConvTranspose3d): + def __init__(self, param): + super(Deconvolution, self).__init__( + first_or(param, "group", 1), + param["num_output"], + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + dilation=first_or(param, "dilation", 1), + groups=first_or(param, "group", 1), + ) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight.numel() == 0 and self.bias.numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(Deconvolution, self).__init__( + x.size(1), + self.out_channels, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + ) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(Deconvolution, self).forward(x) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, weight=weight, bias=bias.view(-1) if bias is not None else bias + ) + self.in_channels = self.weight.size(1) + + +class InnerProduct(nn.Linear): + def __init__(self, param): + super(InnerProduct, self).__init__(1, param["num_output"]) + self.weight, self.bias = nn.Parameter(), nn.Parameter() + self.weight_init, self.bias_init = param.get("weight_filler", {}), param.get( + "bias_filler", {} + ) + + def forward(self, x): + if self.weight.numel() == 0 and self.bias.numel() == 0: + requires_grad = [self.weight.requires_grad, self.bias.requires_grad] + super(InnerProduct, self).__init__(x.size(1), self.out_features) + convert_to_gpu_if_enabled(self) + init_weight_bias(self, requires_grad=requires_grad) + return super(InnerProduct, self).forward( + x if x.size(-1) == self.in_features else x.view(len(x), -1) + ) + + def set_parameters(self, weight=None, bias=None): + init_weight_bias( + self, + weight=weight.view(weight.size(-2), weight.size(-1)) + if weight is not None + else None, + bias=bias.view(-1) if bias is not None else None, + ) + self.in_features = self.weight.size(1) + + +class Concat: + def __call__(self, *args): + return torch.cat(tuple([*args]), dim=1) + + +# using dict calls (=) otherwise the dict does not play nice with +# the lambda function +modules3d = dict( + Convolution=lambda param: Convolution(param), + Deconvolution=lambda param: Deconvolution(param), + InnerProduct=lambda param: InnerProduct(param), + Pooling=lambda param: [nn.MaxPool3d, nn.AvgPool3d][param["pool"]]( + kernel_size=first_or(param, "kernel_size", 1), + stride=first_or(param, "stride", 1), + padding=first_or(param, "pad", 0), + ), + Softmax=lambda param: nn.Softmax(dim=param.get("axis", -1)), + ReLU=lambda param: nn.ReLU(), + Dropout=lambda param: nn.Dropout(p=param["dropout_ratio"]), + Eltwise=lambda param: [torch.mul, torch.add, torch.max][param.get("operation", 1)], + Concat=lambda param: Concat(), + LRN=lambda param: nn.LocalResponseNorm( + size=param["local_size"], alpha=param["alpha"], beta=param["beta"] + ), +) From 1212c180c229c9a5fecc44c3a050631db8bb9ef6 Mon Sep 17 00:00:00 2001 From: dirk Date: Wed, 29 Mar 2023 15:02:26 +0200 Subject: [PATCH 5/7] Add Deconvolution and concatenation layer --- layers2d.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/layers2d.py b/layers2d.py index ea1c3e4..caf961d 100644 --- a/layers2d.py +++ b/layers2d.py @@ -219,7 +219,7 @@ def __init__(self, param): ) def forward(self, x): - if self.weight_numel() == 0 and self.bias_numel() == 0: + if self.weight.numel() == 0 and self.bias.numel() == 0: requires_grad = [self.weight.requires_grad, self.bias.requires_grad] super(Deconvolution, self).__init__( x.size(1), @@ -269,6 +269,11 @@ def set_parameters(self, weight=None, bias=None): self.in_features = self.weight.size(1) +class Concat: + def __call__(self, *args): + return torch.cat(tuple([*args]), dim=1) + + # using dict calls (=) otherwise the dict does not play nice with # the lambda function modules2d = dict( From 5746283ec27ee8b28f7f3c4e3c02b15dd027156e Mon Sep 17 00:00:00 2001 From: dirk Date: Thu, 30 Mar 2023 09:05:31 +0200 Subject: [PATCH 6/7] Add installation instruction --- README.md | 10 ++++++- requirements.txt | 78 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) create mode 100644 requirements.txt diff --git a/README.md b/README.md index a911e7e..f98b1e8 100644 --- a/README.md +++ b/README.md @@ -16,12 +16,20 @@ The layer support isn't as complete as in https://github.com/marvis/pytorch-caff * softmax (axis) * local response norm (local_size, alpha, beta) -Dependencies: protobuf with Python bindings, including `protoc` binary in `PATH`. PRs to enable other layers or layer params are very welcome (see the definition of the `modules` dictionary in the code)! License is MIT. +## Installation Instruction +### Python Environment +This converter is currently working on python3.8 in Ubuntu. To install, create a virtual environment using ```python3 -m vev .venv```. Activate the environment using ```source .venv\bin\activate```. Lastly install the requirements using ```pip install -f requirements.txt```. + +### Protobuf Compiler +Dependencies: protobuf with Python bindings, including `protoc` binary in `PATH`. + +To install protobuf in ubuntu run: ```sudo apt install protobuf-compiler```. + ## Dump weights to PT or HDF5 ```shell # prototxt and caffemodel from https://gist.github.com/ksimonyan/211839e770f7b538e2d8#file-readme-md diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..93d86aa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,78 @@ +astroid==2.15.1 +beautifulsoup4==4.12.0 +black==23.1.0 +cachetools==5.3.0 +certifi==2022.12.7 +charset-normalizer==3.1.0 +click==8.1.3 +cmake==3.26.1 +contourpy==1.0.7 +cycler==0.11.0 +dill==0.3.6 +filelock==3.10.7 +fonttools==4.39.2 +google==3.0.0 +google-api-core==2.11.0 +google-api-python-client==2.82.0 +google-auth==2.16.3 +google-auth-httplib2==0.1.0 +googleapis-common-protos==1.59.0 +h5py==3.8.0 +httplib2==0.22.0 +idna==3.4 +imageio==2.27.0 +importlib-resources==5.12.0 +isort==5.12.0 +Jinja2==3.1.2 +kiwisolver==1.4.4 +lazy-loader==0.2 +lazy-object-proxy==1.9.0 +lit==16.0.0 +MarkupSafe==2.1.2 +matplotlib==3.7.1 +mccabe==0.7.0 +mpmath==1.3.0 +mypy-extensions==1.0.0 +networkx==3.0 +numpy==1.24.2 +nvidia-cublas-cu11==11.10.3.66 +nvidia-cuda-cupti-cu11==11.7.101 +nvidia-cuda-nvrtc-cu11==11.7.99 +nvidia-cuda-runtime-cu11==11.7.99 +nvidia-cudnn-cu11==8.5.0.96 +nvidia-cufft-cu11==10.9.0.58 +nvidia-curand-cu11==10.2.10.91 +nvidia-cusolver-cu11==11.4.0.1 +nvidia-cusparse-cu11==11.7.4.91 +nvidia-nccl-cu11==2.14.3 +nvidia-nvtx-cu11==11.7.91 +packaging==23.0 +pathspec==0.11.1 +Pillow==9.4.0 +platformdirs==3.2.0 +protobuf==4.22.1 +pyasn1==0.4.8 +pyasn1-modules==0.2.8 +pylint==2.17.1 +pyparsing==3.0.9 +python-dateutil==2.8.2 +PyWavelets==1.4.1 +requests==2.28.2 +rsa==4.9 +scikit-image==0.20.0 +scipy==1.9.1 +SimpleITK==2.2.1 +six==1.16.0 +soupsieve==2.4 +sympy==1.11.1 +tifffile==2023.3.21 +tomli==2.0.1 +tomlkit==0.11.7 +torch==2.0.0 +torchsummary==1.5.1 +triton==2.0.0 +typing-extensions==4.5.0 +uritemplate==4.1.1 +urllib3==1.26.15 +wrapt==1.15.0 +zipp==3.15.0 From 4d49dc3cd390c40be70b04c892e40f25777cd07a Mon Sep 17 00:00:00 2001 From: dirk Date: Thu, 30 Mar 2023 09:07:50 +0200 Subject: [PATCH 7/7] Fix typo --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index f98b1e8..1cbab7a 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ License is MIT. ## Installation Instruction ### Python Environment -This converter is currently working on python3.8 in Ubuntu. To install, create a virtual environment using ```python3 -m vev .venv```. Activate the environment using ```source .venv\bin\activate```. Lastly install the requirements using ```pip install -f requirements.txt```. +This converter is currently working on python3.8 in Ubuntu. To install, create a virtual environment using ```python3 -m vev .venv```. Activate the environment using ```source .venv/bin/activate```. Lastly install the requirements using ```pip install -r requirements.txt```. ### Protobuf Compiler Dependencies: protobuf with Python bindings, including `protoc` binary in `PATH`.