model serialization code changed

yunjey · yunjey · commit 3a7a9cf07b91 · 2017-03-18T16:28:19.000+09:00
diff --git a/tutorials/01 - Linear Regression/main.py b/tutorials/01 - Linear Regression/main.py
@@ -61,4 +61,4 @@ def forward(self, x):
 plt.show()
 
 # Save the Model
-torch.save(model, 'model.pkl')
+torch.save(model.state_dict(), 'model.pkl')
diff --git a/tutorials/02 - Logistic Regression/main.py b/tutorials/02 - Logistic Regression/main.py
@@ -79,4 +79,4 @@ def forward(self, x):
 print('Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total))
 
 # Save the Model
-torch.save(model, 'model.pkl')
+torch.save(model.state_dict(), 'model.pkl')
diff --git a/tutorials/03 - Feedforward Neural Network/main-gpu.py b/tutorials/03 - Feedforward Neural Network/main-gpu.py
@@ -81,4 +81,7 @@ def forward(self, x):
     total += labels.size(0)
     correct += (predicted.cpu() == labels).sum()
 
-print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
+print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
+
+# Save the Model
+torch.save(net.state_dict(), 'model.pkl')
diff --git a/tutorials/03 - Feedforward Neural Network/main.py b/tutorials/03 - Feedforward Neural Network/main.py
@@ -81,4 +81,7 @@ def forward(self, x):
     total += labels.size(0)
     correct += (predicted == labels).sum()
 
-print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
+print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
+
+# Save the Model
+torch.save(net.state_dict(), 'model.pkl')
diff --git a/tutorials/04 - Convolutional Neural Network/main-gpu.py b/tutorials/04 - Convolutional Neural Network/main-gpu.py
@@ -90,4 +90,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total))
 
 # Save the Trained Model
-torch.save(cnn, 'cnn.pkl')
+torch.save(cnn.state_dict(), 'cnn.pkl')
diff --git a/tutorials/04 - Convolutional Neural Network/main.py b/tutorials/04 - Convolutional Neural Network/main.py
@@ -90,4 +90,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total))
 
 # Save the Trained Model
-torch.save(cnn, 'cnn.pkl')
+torch.save(cnn.state_dict(), 'cnn.pkl')
diff --git a/tutorials/05 - Deep Residual Network/main-gpu.py b/tutorials/05 - Deep Residual Network/main-gpu.py
@@ -144,4 +144,4 @@ def forward(self, x):
 print('Accuracy of the model on the test images: %d %%' % (100 * correct / total))
 
 # Save the Model
-torch.save(resnet, 'resnet.pkl')
+torch.save(resnet.state_dict(), 'resnet.pkl')
diff --git a/tutorials/05 - Deep Residual Network/main.py b/tutorials/05 - Deep Residual Network/main.py
@@ -103,7 +103,7 @@ def forward(self, x):
         return out
     
 resnet = ResNet(ResidualBlock, [2, 2, 2, 2])
-resnet
+
 
 # Loss and Optimizer
 criterion = nn.CrossEntropyLoss()
@@ -127,7 +127,7 @@ def forward(self, x):
             print ("Epoch [%d/%d], Iter [%d/%d] Loss: %.4f" %(epoch+1, 80, i+1, 500, loss.data[0]))
 
     # Decaying Learning Rate
-    if (epoch+1) % 30 == 0:
+    if (epoch+1) % 20 == 0:
         lr /= 3
         optimizer = torch.optim.Adam(resnet.parameters(), lr=lr) 
 
@@ -144,4 +144,4 @@ def forward(self, x):
 print('Accuracy of the model on the test images: %d %%' % (100 * correct / total))
 
 # Save the Model
-torch.save(resnet, 'resnet.pkl')
+torch.save(resnet.state_dict(), 'resnet.pkl')
diff --git a/tutorials/06 - Recurrent Neural Network/main-gpu.py b/tutorials/06 - Recurrent Neural Network/main-gpu.py
@@ -92,4 +92,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total)) 
 
 # Save the Model
-torch.save(rnn, 'rnn.pkl')
+torch.save(rnn.state_dict(), 'rnn.pkl')
diff --git a/tutorials/06 - Recurrent Neural Network/main.py b/tutorials/06 - Recurrent Neural Network/main.py
@@ -92,4 +92,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total)) 
 
 # Save the Model
-torch.save(rnn, 'rnn.pkl')
+torch.save(rnn.state_dict(), 'rnn.pkl')
diff --git a/tutorials/07 - Bidirectional Recurrent Neural Network/main-gpu.py b/tutorials/07 - Bidirectional Recurrent Neural Network/main-gpu.py
@@ -93,4 +93,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total)) 
 
 # Save the Model
-torch.save(rnn, 'rnn.pkl')
+torch.save(rnn.state_dict(), 'rnn.pkl')
diff --git a/tutorials/07 - Bidirectional Recurrent Neural Network/main.py b/tutorials/07 - Bidirectional Recurrent Neural Network/main.py
@@ -93,4 +93,4 @@ def forward(self, x):
 print('Test Accuracy of the model on the 10000 test images: %d %%' % (100 * correct / total)) 
 
 # Save the Model
-torch.save(rnn, 'rnn.pkl')
+torch.save(rnn.state_dict(), 'rnn.pkl')
diff --git a/tutorials/08 - Language Model/main-gpu.py b/tutorials/08 - Language Model/main-gpu.py
@@ -31,7 +31,6 @@ def __init__(self, vocab_size, embed_size, hidden_size, num_layers):
         self.embed = nn.Embedding(vocab_size, embed_size)
         self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
         self.linear = nn.Linear(hidden_size, vocab_size)
-        
         self.init_weights()
         
     def init_weights(self):
@@ -120,4 +119,4 @@ def detach(states):
             print('Sampled [%d/%d] words and save to %s'%(i+1, num_samples, sample_path))
 
 # Save the Trained Model
-torch.save(model, 'model.pkl')
+torch.save(model.state_dict(), 'model.pkl')
diff --git a/tutorials/08 - Language Model/main.py b/tutorials/08 - Language Model/main.py
@@ -31,7 +31,6 @@ def __init__(self, vocab_size, embed_size, hidden_size, num_layers):
         self.embed = nn.Embedding(vocab_size, embed_size)
         self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
         self.linear = nn.Linear(hidden_size, vocab_size)
-        
         self.init_weights()
         
     def init_weights(self):
@@ -120,4 +119,4 @@ def detach(states):
             print('Sampled [%d/%d] words and save to %s'%(i+1, num_samples, sample_path))
 
 # Save the Trained Model
-torch.save(model, 'model.pkl')
+torch.save(model.state_dict(), 'model.pkl')
diff --git a/tutorials/10 - Generative Adversarial Network/main-gpu.py b/tutorials/10 - Generative Adversarial Network/main-gpu.py
@@ -23,8 +23,8 @@
                                            batch_size=100, 
                                            shuffle=True)
 
-# 5x5 Convolution
-def conv5x5(in_channels, out_channels, stride):
+# 4x4 Convolution
+def conv4x4(in_channels, out_channels, stride):
     return nn.Conv2d(in_channels, out_channels, kernel_size=4, 
                      stride=stride, padding=1, bias=False)
 
@@ -33,12 +33,12 @@ class Discriminator(nn.Module):
     def __init__(self):
         super(Discriminator, self).__init__()
         self.model = nn.Sequential(
-            conv5x5(3, 16, 2),
+            conv4x4(3, 16, 2),
             nn.LeakyReLU(0.2, inplace=True),
-            conv5x5(16, 32, 2),
+            conv4x4(16, 32, 2),
             nn.BatchNorm2d(32),
             nn.LeakyReLU(0.2, inplace=True),
-            conv5x5(32, 64, 2), 
+            conv4x4(32, 64, 2), 
             nn.BatchNorm2d(64),
             nn.LeakyReLU(0.2, inplace=True),
             nn.Conv2d(64, 1, kernel_size=4),
@@ -91,16 +91,16 @@ def forward(self, x):
 for epoch in range(50):
     for i, (images, _) in enumerate(train_loader):
         images = Variable(images.cuda())
-        real_labels = Variable(torch.ones(images.size(0)).cuda())
-        fake_labels = Variable(torch.zeros(images.size(0)).cuda())
+        real_labels = Variable(torch.ones(images.size(0))).cuda()
+        fake_labels = Variable(torch.zeros(images.size(0))).cuda()
         
         # Train the discriminator
         discriminator.zero_grad()
         outputs = discriminator(images)
         real_loss = criterion(outputs, real_labels)
         real_score = outputs
         
-        noise = Variable(torch.randn(images.size(0), 128).cuda())
+        noise = Variable(torch.randn(images.size(0), 128)).cuda()
         fake_images = generator(noise)
         outputs = discriminator(fake_images) 
         fake_loss = criterion(outputs, fake_labels)
@@ -112,7 +112,7 @@ def forward(self, x):
         
         # Train the generator 
         generator.zero_grad()
-        noise = Variable(torch.randn(images.size(0), 128).cuda())
+        noise = Variable(torch.randn(images.size(0), 128)).cuda()
         fake_images = generator(noise)
         outputs = discriminator(fake_images)
         g_loss = criterion(outputs, real_labels)
@@ -130,5 +130,5 @@ def forward(self, x):
                 './data/fake_samples_%d_%d.png' %(epoch+1, i+1))
 
 # Save the Models 
-torch.save(generator, './generator.pkl')
-torch.save(discriminator, './discriminator.pkl')
+torch.save(generator.state_dict(), './generator.pkl')
+torch.save(discriminator.state_dict(), './discriminator.pkl')
diff --git a/tutorials/10 - Generative Adversarial Network/main.py b/tutorials/10 - Generative Adversarial Network/main.py
@@ -23,8 +23,8 @@
                                            batch_size=100, 
                                            shuffle=True)
 
-# 5x5 Convolution
-def conv5x5(in_channels, out_channels, stride):
+# 4x4 Convolution
+def conv4x4(in_channels, out_channels, stride):
     return nn.Conv2d(in_channels, out_channels, kernel_size=4, 
                      stride=stride, padding=1, bias=False)
 
@@ -33,12 +33,12 @@ class Discriminator(nn.Module):
     def __init__(self):
         super(Discriminator, self).__init__()
         self.model = nn.Sequential(
-            conv5x5(3, 16, 2),
+            conv4x4(3, 16, 2),
             nn.LeakyReLU(0.2, inplace=True),
-            conv5x5(16, 32, 2),
+            conv4x4(16, 32, 2),
             nn.BatchNorm2d(32),
             nn.LeakyReLU(0.2, inplace=True),
-            conv5x5(32, 64, 2), 
+            conv4x4(32, 64, 2), 
             nn.BatchNorm2d(64),
             nn.LeakyReLU(0.2, inplace=True),
             nn.Conv2d(64, 1, kernel_size=4),
@@ -90,7 +90,7 @@ def forward(self, x):
 # Training 
 for epoch in range(50):
     for i, (images, _) in enumerate(train_loader):
-        images = Variable(images.cuda())
+        images = Variable(images)
         real_labels = Variable(torch.ones(images.size(0)))
         fake_labels = Variable(torch.zeros(images.size(0)))
         
@@ -130,5 +130,5 @@ def forward(self, x):
                 './data/fake_samples_%d_%d.png' %(epoch+1, i+1))
 
 # Save the Models
-torch.save(generator, './generator.pkl')
-torch.save(discriminator, './discriminator.pkl')
+torch.save(generator.state_dict(), './generator.pkl')
+torch.save(discriminator.state_dict(), './discriminator.pkl')
