Hello! I keep getting an error when attempting to save my CNN model. I do not have this problem for smaller images like the 32x32 cifar dataset; however, my images are 448x672 (note: a multiple of 224). I am using the model for a regression task. Any help would be much appreciated!!
python 3.7.5
pytorch 1.6.0
Anaconda
Here is my model:
class Network_CNN_batchNorm(nn.Module):
def __init__(self):
super(Network_CNN_batchNorm, self).__init__()
# 3x448x672 input image (RGB)
self.layer1 = nn.Sequential(
# input is 3 channels (RGB) - first parameter
# 64 filters of kernel size 5x5; padding = kernel_size/2 - 1;
nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
# max pooling with stride=2 makes output image 224x336
nn.MaxPool2d(kernel_size=2, stride=2),
nn.BatchNorm2d(64))
self.layer2 = nn.Sequential(
# 2nd layer uses 128 channels (filters) of 3x3
nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.BatchNorm2d(128))
# 3rd layer uses 128 channels (filters) of 3x3
# output feature map is still 224x336
self.layer3 = nn.Sequential(
nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.BatchNorm2d(128))
# Average Pooling Layer, 112x168 output
self.avgP1 = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
# Fully connected layers
self.fc1 = nn.Linear(112 * 168 * 128, 1000)
self.fc2 = nn.Linear(1000, 10) # 10 outputs
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = self.avgP1(out)
out = out.reshape(out.size(0), -1) # flatten
out = self.fc1(out)
out = self.fc2(out)
return out
Note that I can see my training and validation loss decrease over multiple epochs so training the model appears to be fine. However, I do notice that when training on my CPU the memory usage is around 30-40Gb which seems excessive.
The code for saving the model is shown, and I can confirm that the path is OK since it works with smaller image sizes.
torch.save(model.state_dict(), os.path.join(Model_Path, 'epoch-{}.pth'.format(epoch)))
The error I am getting is as follows:
File "C:\my.py", line 526, in <module>
model_trained, t_loss, v_loss = train_model(model, criterion, optimizer, trainloader, testloader, num_epochs)
File "C:\my.py", line 356, in train_model
torch.save(model.state_dict(), os.path.join(Model_Path, 'epoch-{}.pth'.format(epoch)))
File "C:\Users\...\anaconda3\envs\TF2.0\lib\site-packages\torch\serialization.py", line 364, in save
_save(obj, opened_zipfile, pickle_module, pickle_protocol)
File "C:\Users\...\anaconda3\envs\TF2.0\lib\site-packages\torch\serialization.py", line 477, in _save
zip_file.write_record(name, storage.data_ptr(), num_bytes)
TypeError: write_record(): incompatible function arguments. The following argument types are supported:
1. (self: torch._C.PyTorchFileWriter, arg0: str, arg1: str, arg2: int) -> None
2. (self: torch._C.PyTorchFileWriter, arg0: str, arg1: int, arg2: int) -> None
Invoked with: <torch._C.PyTorchFileWriter object at 0x0000026AD2154D30>, 'data/2657683100064', 2657910136960, -7546077184
