| from utils import CustomDataset, transform, preproc, Convert_ONNX |
| from torch.utils.data import Dataset, DataLoader |
| import torch |
| import numpy as np |
| from resnet_model import ResidualBlock, ResNet |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| import tqdm |
| import torch.nn.functional as F |
| from torch.optim.lr_scheduler import ReduceLROnPlateau |
| import pickle |
| import sys |
|
|
| ind = int(sys.argv[1]) |
| seeds = [1,42,7109,2002,32] |
| seed = seeds[ind] |
| print("using seed: ",seed) |
| torch.manual_seed(seed) |
|
|
|
|
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
| num_gpus = torch.cuda.device_count() |
| print(num_gpus) |
|
|
| |
| data_dir = '/mnt/buf1/pma/frbnn/train_ready' |
| dataset = CustomDataset(data_dir, transform=transform) |
| valid_data_dir = '/mnt/buf1/pma/frbnn/valid_ready' |
| valid_dataset = CustomDataset(valid_data_dir, transform=transform) |
|
|
|
|
| num_classes = 2 |
| trainloader = DataLoader(dataset, batch_size=420, shuffle=True, num_workers=32) |
| validloader = DataLoader(valid_dataset, batch_size=512, shuffle=True, num_workers=32) |
|
|
| model = ResNet(24, ResidualBlock, [3, 4, 6, 3], num_classes=num_classes).to(device) |
| model = nn.DataParallel(model) |
| model = model.to(device) |
| params = sum(p.numel() for p in model.parameters()) |
| print("num params ",params) |
| torch.save(model.state_dict(), 'models/test.pt') |
| model.load_state_dict(torch.load('models/test.pt')) |
|
|
| preproc_model = preproc() |
| Convert_ONNX(model.module,'models/test.onnx', input_data_mock=torch.randn(1, 24, 192, 256).to(device)) |
| Convert_ONNX(preproc_model,'models/preproc.onnx', input_data_mock=torch.randn(32, 192, 2048).to(device)) |
|
|
| |
|
|
| criterion = nn.CrossEntropyLoss(weight = torch.tensor([1,1]).to(device)) |
| optimizer = optim.Adam(model.parameters(), lr=0.0001) |
| scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10) |
|
|
|
|
| from tqdm import tqdm |
| |
| epochs = 10000 |
| for epoch in range(epochs): |
| running_loss = 0.0 |
| correct_train = 0 |
| total_train = 0 |
| with tqdm(trainloader, unit="batch") as tepoch: |
| model.train() |
| for i, (images, labels) in enumerate(tepoch): |
| inputs, labels = images.to(device), labels.to(device).float() |
| optimizer.zero_grad() |
| outputs = model(inputs, return_mask=False).to(device) |
| new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32).to(device) |
| loss = criterion(outputs, new_label) |
| loss.backward() |
| optimizer.step() |
| running_loss += loss.item() |
| |
| _, predicted = torch.max(outputs.data, 1) |
| total_train += labels.size(0) |
| correct_train += (predicted == labels).sum().item() |
| val_loss = 0.0 |
| correct_valid = 0 |
| total = 0 |
| model.eval() |
| with torch.no_grad(): |
| for images, labels in validloader: |
| inputs, labels = images.to(device), labels.to(device).float() |
| optimizer.zero_grad() |
| outputs = model(inputs, return_mask=False) |
| new_label = F.one_hot(labels.type(torch.int64),num_classes=2).type(torch.float32) |
| loss = criterion(outputs, new_label) |
| val_loss += loss.item() |
| _, predicted = torch.max(outputs, 1) |
| total += labels.size(0) |
| correct_valid += (predicted == labels).sum().item() |
| scheduler.step(val_loss) |
| |
| train_accuracy = 100 * correct_train / total_train |
| val_accuracy = correct_valid / total * 100.0 |
| torch.save(model.state_dict(), 'models/model-'+str(epoch)+'-'+str(val_accuracy)+'.pt') |
| Convert_ONNX(model.module,'models/model-'+str(epoch)+'-'+str(val_accuracy)+'.onnx', input_data_mock=inputs) |
|
|
| print("===========================") |
| print('accuracy: ', epoch, train_accuracy, val_accuracy) |
| print('learning rate: ', scheduler.get_last_lr()) |
| print("===========================") |
| if scheduler.get_last_lr()[0] <1e-6: |
| break |
|
|
