(item, inputs, input_shape, args, kwargs={})
| 4345 | return super().is_leaf_module(m, module_qualified_name) |
| 4346 | |
| 4347 | def process(item, inputs, input_shape, args, kwargs={}): |
| 4348 | # Skip assertion and validation functions from torch.fx trace |
| 4349 | if callable(item) and hasattr(item, '__name__') and ( |
| 4350 | item.__name__.startswith('_assert') or |
| 4351 | item.__name__ in ('eq', 'getitem', 'size')): |
| 4352 | return |
| 4353 | if item == torch.cat: |
| 4354 | if len(inputs) > 1: |
| 4355 | layers.append( |
| 4356 | Concat(inputs, dimension=len(inputs[0].shape) - 1)) |
| 4357 | return |
| 4358 | elif item == operator.add: |
| 4359 | layers.append(Add(inputs)) |
| 4360 | return |
| 4361 | elif item in (torch.flatten, 'flatten', 'size'): |
| 4362 | return |
| 4363 | elif item == 'view': |
| 4364 | assert -1 in args or \ |
| 4365 | reduce(operator.mul, args) == reduce(operator.mul, input_shape) |
| 4366 | return |
| 4367 | elif item == torch.nn.functional.avg_pool2d: |
| 4368 | layers.append(FixAveragePool2d(input_shape, None, args[1], |
| 4369 | kwargs.get('stride', args[1]), |
| 4370 | kwargs.get('padding', 0))) |
| 4371 | input_shape = layers[-1].shape |
| 4372 | return |
| 4373 | # single-input layers from here |
| 4374 | if inputs and len(inputs) > 1: |
| 4375 | raise CompilerError('multi-input layer %s not supported' % item) |
| 4376 | name = type(item).__name__ |
| 4377 | if name == 'Linear': |
| 4378 | # Precondition: the item |
| 4379 | assert item.bias is not None |
| 4380 | if mul(input_shape[1:]) == item.in_features: |
| 4381 | layers.append(Dense(input_shape[0], item.in_features, |
| 4382 | item.out_features)) |
| 4383 | elif input_shape[-1] == item.in_features: |
| 4384 | # we loop over all but last dimension |
| 4385 | assert len(input_shape) == 3, "Dense only supports one extra dimension to loop over" |
| 4386 | d = input_shape[1] |
| 4387 | layers.append(Dense(input_shape[0], item.in_features, |
| 4388 | item.out_features, d)) |
| 4389 | else: |
| 4390 | assert False, f"input shape {input_shape} incompatible with in_features {item.in_features}" |
| 4391 | |
| 4392 | if input_via is not None: |
| 4393 | shapes = [x.shape for x in (layers[-1].W, layers[-1].b)] |
| 4394 | import numpy |
| 4395 | swapped = item.weight.detach().numpy() |
| 4396 | if len(input_shape) == 4: |
| 4397 | print (swapped.shape) |
| 4398 | swapped = numpy.reshape( |
| 4399 | swapped, |
| 4400 | [item.out_features, input_shape[3]] + input_shape[1:3]) |
| 4401 | print (swapped.shape) |
| 4402 | swapped = numpy.moveaxis(swapped, 1, -1) |
| 4403 | print (swapped.shape) |
| 4404 | swapped = numpy.reshape( |
no test coverage detected