| import os |
| import time |
|
|
| from data.dataset import TextDataset |
| from models.model import WriteViT
|
| from params import *
|
|
|
| def main():
|
|
|
| init_project() |
|
|
| TextDatasetObj = TextDataset(num_examples = NUM_EXAMPLES)
|
| dataset = torch.utils.data.DataLoader(
|
| TextDatasetObj,
|
| batch_size=batch_size,
|
| shuffle=True,
|
| num_workers=0,
|
| pin_memory=True, drop_last=True,
|
| collate_fn=TextDatasetObj.collate_fn)
|
|
|
| model = WriteViT(backbone=BACKBONE).to(DEVICE) |
|
|
| os.makedirs('saved_models', exist_ok = True)
|
| MODEL_PATH = os.path.join('saved_models', EXP_NAME)
|
| if os.path.isdir(MODEL_PATH) and RESUME:
|
| model.load_state_dict(torch.load(MODEL_PATH+'/model.pth'))
|
| print (MODEL_PATH+' : Model loaded Successfully')
|
| else:
|
| if not os.path.isdir(MODEL_PATH): os.mkdir(MODEL_PATH)
|
|
|
|
|
| for epoch in range(EPOCHS):
|
|
|
|
|
| start_time = time.time()
|
|
|
| for i,data in enumerate(dataset):
|
|
|
| if (i % NUM_CRITIC_GOCR_TRAIN) == 0:
|
|
|
| model._set_input(data)
|
| model.optimize_G_only()
|
| model.optimize_G_step()
|
|
|
| if (i % NUM_CRITIC_DOCR_TRAIN) == 0:
|
|
|
| model._set_input(data)
|
| model.optimize_D_OCR_W()
|
| model.optimize_D_OCR_W_step()
|
|
|
|
|
|
|
| end_time = time.time()
|
|
|
| losses = model.get_current_losses()
|
|
|
| print ({'EPOCH':epoch, 'TIME':end_time-start_time, 'LOSSES': losses}) |
|
|
| if epoch % SAVE_MODEL == 0: torch.save(model.state_dict(), MODEL_PATH+ '/model.pth')
|
| if epoch % SAVE_MODEL_HISTORY == 0: torch.save(model.state_dict(), MODEL_PATH+ '/model'+str(epoch)+'.pth')
|
|
|
|
|
| if __name__ == "__main__":
|
|
|
| main() |
|
|
