diff options
| author | ericmarin <maarin.eric@gmail.com> | 2026-03-26 20:28:38 +0100 |
|---|---|---|
| committer | ericmarin <maarin.eric@gmail.com> | 2026-03-26 21:27:10 +0100 |
| commit | 3e338c3be65638ef1898c32c707c50422acafb18 (patch) | |
| tree | 80a29ac6b7baee3bbfe4f161fc893fd5948d9409 /fashion_mnist/fashion_mnist.py | |
| parent | 689c34076d08e59b1382864f9efcd983c8665ae5 (diff) | |
| download | vein-3e338c3be65638ef1898c32c707c50422acafb18.tar.gz vein-3e338c3be65638ef1898c32c707c50422acafb18.zip | |
added LICENSE
Diffstat (limited to 'fashion_mnist/fashion_mnist.py')
| -rw-r--r-- | fashion_mnist/fashion_mnist.py | 53 |
1 files changed, 0 insertions, 53 deletions
diff --git a/fashion_mnist/fashion_mnist.py b/fashion_mnist/fashion_mnist.py deleted file mode 100644 index 1c9dcf7..0000000 --- a/fashion_mnist/fashion_mnist.py +++ /dev/null @@ -1,53 +0,0 @@ -import torch, torch.nn as nn -from torchvision.datasets import FashionMNIST -from torch.utils.data import DataLoader -from torchvision import transforms - -class FashionMNIST_MLP(nn.Module): - def __init__(self, hidden_dim): - super().__init__() - self.layers = nn.Sequential( - nn.Flatten(), - nn.Linear(6 * 6, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, 2) - ) - def forward(self, x): - return self.layers(x) - -transform = transforms.Compose([ - transforms.Resize((6, 6)), - transforms.ToTensor(), -]) - -train_dataset = FashionMNIST('./', download=True, transform=transform, train=True) -tshirts_trousers = [id for id, data in enumerate(train_dataset.targets) if data.item() == 0 or data.item() == 1] -train_dataset = torch.utils.data.Subset(train_dataset, tshirts_trousers) - -trainloader = DataLoader(train_dataset, batch_size=128, shuffle=True) - -def train_model(name: str, dim): - net = FashionMNIST_MLP(hidden_dim=dim) - loss_fn = nn.CrossEntropyLoss() - optimizer = torch.optim.Adam(net.parameters(), lr=1e-4) - - print(f"Training {name}...") - for epoch in range(100): - global loss - for data in trainloader: - inputs, targets = data - optimizer.zero_grad() - outputs = net(inputs) - loss = loss_fn(outputs, targets) - loss.backward() - optimizer.step() - if (epoch + 1) % 10 == 0: - print(f" Epoch {epoch+1}, Loss: {loss.item():.4f}") - return net - -if __name__ == "__main__": - torch_net_a = train_model("Network A", 6).eval() - torch_net_b = train_model("Network B", 12).eval() - - torch.onnx.export(torch_net_a, (torch.randn(1, 1, 6, 6),), "fashion_mnist_a.onnx") - torch.onnx.export(torch_net_b, (torch.randn(1, 1, 6, 6),), "fashion_mnist_b.onnx") |
