119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
import time
|
|
import numpy as np
|
|
import torch
|
|
from torch.nn.functional import *
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from torch import nn
|
|
from torchvision import datasets, transforms
|
|
from tqdm import tqdm
|
|
from utils import *
|
|
|
|
import ipdb
|
|
|
|
|
|
class Model_2_1(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.fc = nn.Linear(in_features=200, out_features=1)
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
x = self.fc(x)
|
|
x = self.sigmoid(x)
|
|
return x
|
|
|
|
class My_BinaryCLS_Dataset(Dataset):
|
|
def __init__(self, train=True, num_features=200):
|
|
num_samples = 7000 if train else 3000
|
|
|
|
x_1 = np.random.normal(loc=-0.5, scale=0.2, size=(num_samples, num_features))
|
|
x_2 = np.random.normal(loc=0.5, scale=0.2, size=(num_samples, num_features))
|
|
|
|
labels_1 = np.zeros((num_samples, 1))
|
|
labels_2 = np.ones((num_samples, 1))
|
|
|
|
x = np.concatenate((x_1, x_2), axis=0)
|
|
labels = np.concatenate((labels_1, labels_2), axis=0)
|
|
self.data = [[x[i], labels[i]] for i in range(2 * num_samples)]
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, index):
|
|
x, y = self.data[index]
|
|
x = torch.FloatTensor(x)
|
|
y = torch.LongTensor(y)
|
|
return x, y
|
|
|
|
|
|
if __name__ == "__main__":
|
|
learning_rate = 1e-4
|
|
num_epochs = 10
|
|
batch_size = 512
|
|
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
|
|
train_binarycls_dataset = My_BinaryCLS_Dataset(train=True)
|
|
test_binarycls_dataset = My_BinaryCLS_Dataset(train=False)
|
|
train_dataloader = DataLoader(
|
|
dataset=train_binarycls_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=14,
|
|
pin_memory=True,
|
|
)
|
|
test_dataloader = DataLoader(
|
|
dataset=test_binarycls_dataset,
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=14,
|
|
pin_memory=True,
|
|
)
|
|
|
|
model = Model_2_1().to(device)
|
|
criterion = nn.BCELoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
|
|
|
|
for epoch in range(num_epochs):
|
|
model.train()
|
|
total_epoch_loss = 0
|
|
start_time = time.time()
|
|
for index, (x, targets) in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
|
|
optimizer.zero_grad()
|
|
|
|
x = x.to(device)
|
|
targets = targets.to(device).to(dtype=torch.float32)
|
|
|
|
y_pred = model(x)
|
|
loss = criterion(y_pred, targets)
|
|
total_epoch_loss += loss.item()
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
end_time = time.time()
|
|
train_time = end_time - start_time
|
|
|
|
model.eval()
|
|
with torch.no_grad():
|
|
total_epoch_acc = 0
|
|
start_time = time.time()
|
|
for index, (x, targets) in tqdm(enumerate(test_dataloader), total=len(test_dataloader)):
|
|
x = x.to(device)
|
|
targets = targets.to(device)
|
|
|
|
output = model(x)
|
|
pred = (output > 0.5).to(dtype=torch.long)
|
|
total_epoch_acc += (pred == targets).sum().item()
|
|
|
|
end_time = time.time()
|
|
test_time = end_time - start_time
|
|
|
|
avg_epoch_acc = total_epoch_acc / len(test_binarycls_dataset)
|
|
print(
|
|
f"Epoch [{epoch + 1}/{num_epochs}],",
|
|
f"Train Loss: {total_epoch_loss},",
|
|
f"Used Time: {train_time * 1000:.3f}ms,",
|
|
f"Test Acc: {avg_epoch_acc * 100:.3f}%,",
|
|
f"Used Time: {test_time * 1000:.3f}ms",
|
|
)
|