self.fc = nn.Sequential(nn.Linear(128 * 4 * 4, 1024),
nn.ReLU(inplace=True), nn.Linear(1024, 128),
nn.ReLU(inplace=True), nn.Linear(128, 10))
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
#cnn = CNN()
cnn = CNN_model.MNISTResNet()
print(cnn)
if torch.cuda.is_available():
cnn.cuda()
# Loss and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learning_rate)
# Train the Model
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
if torch.cuda.is_available():
images = Variable(images).cuda()
labels = Variable(labels).cuda()
else:
def __init__(self):
self.writeimg = True
self.model = CNN_model.MNISTResNet()
self.model.load_state_dict(torch.load('Res.pth'))
self.model.eval()