def _load(self, filePath): checkpoint = torch.load(filePath) model = MyModel( device, checkpoint['inputSize'], checkpoint['gatedCnnOutputSize'], checkpoint['gatedCnnStride1'], checkpoint['gatedCnnStride2'], checkpoint['gatedCnnKernel1'], checkpoint['gatedCnnKernel2'], checkpoint['lstmLayer'], checkpoint['lstmHiddenSize'], checkpoint['fcOutputSize'], checkpoint['dropout']) model.load_state_dict(checkpoint['stateDict']) model.eval() if self.device.type == 'cpu': model.cpu() else: model.cuda(device=self.device) return model
dtype=torch.float32).to(device)) #criterion = MyLoss2(weights = torch.tensor(train_dataset.get_class_weights2(), dtype=torch.float32).to(device)) optimizer = optim.SGD(model.parameters(), lr=LR) #scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=2, factor = 0.5) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1) #scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10,15,20,25,30], gamma=0.1) train_dataloader = DataLoader(train_dataset,\ batch_size = BATCHSIZE,\ shuffle = True, num_workers=4 ) validation_dataloader = DataLoader(validation_dataset,\ batch_size = BATCHSIZE,\ shuffle = False ) pre_score = 0 for epoch in range(EPOCHS): train(model, train_dataloader, criterion, optimizer, device) if (epoch + 1) % 5 == 0: score = validation(model, validation_dataloader, criterion, None, device) if score > pre_score: pre_score = score model = model.cpu() torch.save(model.state_dict(), save_path) model = model.to(device) scheduler.step(score)