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()