| 43 | |
| 44 | import torch |
| 45 | class MyModel(torch.nn.Module): |
| 46 | def __init__(self) -> None: |
| 47 | super().__init__() |
| 48 | with torch.no_grad(): |
| 49 | self.conv1 = torch.nn.Conv2d(in_channels=8, out_channels=32, kernel_size=3, stride=2, padding=1) |
| 50 | self.conv2 = torch.nn.Conv2d(in_channels=32, out_channels=32, groups=16, kernel_size=3, stride=1, padding=1, bias=False) |
| 51 | self.conv3 = torch.nn.Conv2d(in_channels=32, out_channels=8, groups=8, kernel_size=5, stride=1, padding=2) |
| 52 | self.convtranspose1 = torch.nn.ConvTranspose2d(in_channels=8, out_channels=32, kernel_size=5, stride=1, padding=2) |
| 53 | self.convtranspose2 = torch.nn.ConvTranspose2d(in_channels=32, out_channels=32, groups=32, kernel_size=3, stride=2, bias=False) |
| 54 | self.convtranspose3 = torch.nn.ConvTranspose2d(in_channels=32, out_channels=8, groups=1, kernel_size=1) |
| 55 | |
| 56 | self.conv1.bias.copy_(torch.rand_like(self.conv1.bias)) |
| 57 | self.conv3.bias.copy_(torch.rand_like(self.conv3.bias)) |
| 58 | self.convtranspose1.bias.copy_(torch.rand_like(self.convtranspose1.bias)) |
| 59 | self.convtranspose3.bias.copy_(torch.rand_like(self.convtranspose3.bias)) |
| 60 | |
| 61 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 62 | x = self.conv1(x) |
| 63 | x = self.conv2(x) |
| 64 | x = self.conv3(x) |
| 65 | x = self.convtranspose1(x) |
| 66 | x = self.convtranspose2(x) |
| 67 | x = self.convtranspose3(x) |
| 68 | return x |
| 69 | |
| 70 | model = MyModel().cuda() |
| 71 | dump_torch_to_onnx(model=model, onnx_export_file='model.onnx', input_shape=[1, 8, 96, 96]) |
no outgoing calls
no test coverage detected