-
-
Save cbaziotis/06392434d4bbb740985e06333203d423 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
from torch.autograd import Variable | |
import torch.nn.functional as F | |
import matplotlib.pyplot as plt | |
import numpy as np | |
torch.manual_seed(2809) | |
np.random.seed(2809) | |
# Create simple model | |
class MyModel(nn.Module): | |
def __init__(self, in_features, hidden, out_features): | |
super(MyModel, self).__init__() | |
self.fc1 = nn.Linear(in_features, hidden) | |
self.drop1 = nn.Dropout(p=0.5) | |
self.bn1 = nn.BatchNorm1d(hidden) | |
self.fc2 = nn.Linear(hidden, out_features) | |
def forward(self, x): | |
x = F.relu(self.fc1(x)) | |
x = self.drop1(x) | |
x = self.bn1(x) | |
x = F.log_softmax(self.fc2(x), dim=1) | |
return x | |
def weight_init(m): | |
if isinstance(m, nn.Linear): | |
m.weight.data.fill_(1.0) | |
m.bias.data.zero_() | |
def train(model, optimizer): | |
optimizer.zero_grad() | |
output = model(x) | |
loss = criterion(output, y) | |
loss.backward() | |
optimizer.step() | |
return loss.data[0] | |
def print_weights(model): | |
print(model.fc1.weight.data) | |
print(model.fc2.weight.data) | |
def save_checkpoint(state, filename='checkpoint.pth.tar'): | |
torch.save(state, filename) | |
checkpoint_path = 'checkpoint.pth.tar' | |
batch_size = 10 | |
in_features = 5 | |
hidden = 10 | |
out_features = 3 | |
model = MyModel(in_features, hidden, out_features) | |
model.apply(weight_init) | |
print_weights(model) | |
optimizer = optim.Adam(model.parameters(), | |
lr=5e-4, | |
betas=(0.9, 0.999), | |
eps=1e-08, | |
weight_decay=1e-4) | |
criterion = nn.NLLLoss() | |
x = Variable(torch.randn(batch_size, in_features)) | |
y = Variable(torch.LongTensor(batch_size).random_(out_features)) | |
# Train for one epoch | |
losses = [] | |
for epoch in range(10): | |
loss_val = train(model, optimizer) | |
losses.append(loss_val) | |
print_weights(model) | |
losses_arr = np.array(losses) | |
plt.plot(losses_arr) | |
save_checkpoint({ | |
'epoch': epoch + 1, | |
'state_dict': model.state_dict(), | |
'optimizer' : optimizer.state_dict(), | |
}, checkpoint_path) | |
# Resume training | |
model_resume = MyModel(in_features, hidden, out_features) | |
optimizer_resume = optim.Adam(model_resume.parameters(), | |
lr=5e-4, | |
betas=(0.9, 0.999), | |
eps=1e-08, | |
weight_decay=1e-4) | |
print("=> loading checkpoint '{}'".format(checkpoint_path)) | |
checkpoint = torch.load(checkpoint_path) | |
start_epoch = checkpoint['epoch'] | |
model_resume.load_state_dict(checkpoint['state_dict']) | |
optimizer_resume.load_state_dict(checkpoint['optimizer']) | |
print("=> loaded checkpoint '{}' (epoch {})" | |
.format(checkpoint_path, checkpoint['epoch'])) | |
print_weights(model) | |
print_weights(model_resume) | |
# Train models for 10 epochs | |
losses_resume = losses | |
for epoch in range(10, 20): | |
loss_val = train(model, optimizer) | |
losses.append(loss_val) | |
loss_val_resume = train(model_resume, optimizer_resume) | |
losses_resume.append(loss_val_resume) | |
losses_arr = np.array(losses) | |
losses_resume_arr = np.array(losses_resume) | |
plt.plot(losses_arr) | |
plt.plot(losses_resume_arr) | |
if not np.allclose(losses_arr, losses_resume_arr): | |
print('Losses are different!') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment