Hello everyone!
I am trying to implement the function inside the model, however, no progress so far…
I have a function which takes three arguments, one of them is an input (image transformed to PyTorch tensor) and two more numbers and returns processed image transformed to PyTorch tensor:
def process_image(image_in, variable_1, variable_2):
img = image_in.cpu().detach().numpy()
#do whatever to get img_out
img_out = transforms.ToTensor()(img_out)
return img_out
then I am trying to add the function in the simple NN:
class MyNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, 1, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, 1, padding=1)
self.conv3 = nn.Conv2d(32, 64, 3, 1, padding=1)
self.fc1 = nn.Linear(1024, 300)
self.dropout1 = nn.Dropout(0.5)
self.fc2 = nn.Linear(300, 10)
def forward(self, x):
x = process_image(x, variable_1, variable_2)
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv3(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 1024)
x = F.relu(self.fc1(x))
x = self.dropout1(x)
x = self.fc2(x)
return x
I guess I am doing it wrong, because, so far, Colab just crashes without any particular explanation…
Also, I’d like to make variable_1
and variable_2
trainable. How can I do that?
Thanks in advance for helping a newbie!