Skip to content

Commit 941b3e1

Browse files
authored
BatchNorm training and cat_breeds readme (Iainmon#57)
BatchNorm doesn't actually work because of type incompatibility: real(32) <-> real(64). This seems fine, I am going to merge.
2 parents 0816102 + da53057 commit 941b3e1

8 files changed

Lines changed: 346 additions & 48 deletions

File tree

examples/cat_breeds/README.md

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
# Cat breeds tutorial
2+
3+
Download the images from https://www.kaggle.com/datasets/imbikramsaha/cat-breeds/code, or with the following code.
4+
5+
```
6+
kagglehub.dataset_download("imbikramsaha/cat-breeds")
7+
```
8+
9+
This tutorial demonstrates training and loading a model and its data into ChAI code to be used for multi-locale inference.
10+
11+
## Image Preprocessing (load_cats.py)
12+
13+
The structure of the data will vary between different Kaggle datasets. For this specific one, the data consists of .jpg images of 12 different cat breeds, with each breed separated into its own directory. The following code iterates through every image in this structure.
14+
15+
```
16+
classes=sorted(os.listdir(sdir) )
17+
n = 0
18+
for i, c in enumerate(classes):
19+
cpath=os.path.join(sdir, c)
20+
files=os.listdir(cpath)
21+
for f in files:
22+
fpath=os.path.join(cpath,f)
23+
```
24+
25+
PyTorch expects images to have (C, H, W) dimensions, which stands for channel, height, and width. It is also much easier to work with images that have the same height and width as each other. Within the for-loops above, every image in the dataset is resized down to (32, 32), transposed from (H, W, C) to (C, H, W), and saved as a .npy file. Labels are also saved as .npy files.
26+
27+
```
28+
image = cv2.imread(fpath)
29+
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
30+
resized_img = cv2.resize(
31+
src=image,
32+
dsize(32, 32),
33+
interpolation=cv2.INTER_CUBIC
34+
)
35+
transposed_img = np.transpose(resized_img, (2, 0, 1))
36+
np.save(f"{save_path}/images/item{n}", transposed_img)
37+
np.save(f"{save_path}/labels/item{n}", i)
38+
```
39+
40+
## Building the Model (models/for_cats.py)
41+
42+
To create a customized model in PyTorch, we create a class that inherits from nn.Module and provides its own __init__() and forward() functions. Define the layers of the model in __init__() and specify how the data will pass through the model in forward().
43+
44+
Although it is not required, using nn.Sequential helps to ensure that every layer or activation function of the model is readable by ChAI.
45+
46+
```
47+
import torch
48+
import torch.nn as nn
49+
import torch.nn.functional as F
50+
51+
class SmallCNN(nn.Module):
52+
def __init__(self):
53+
super(SmallCNN, self).__init__()
54+
self.layers = nn.Sequential(
55+
nn.Conv2d(3, 64, 3, padding="same"),
56+
nn.ReLU(),
57+
nn.Conv2d(64, 128, 3, padding="same"),
58+
nn.ReLU(),
59+
nn.MaxPool2d(2, 2),
60+
nn.Flatten(),
61+
nn.Linear(8192, 256),
62+
nn.Linear(256, 10)
63+
)
64+
65+
def forward(self, x):
66+
return self.layers(x)
67+
```
68+
69+
## Loading Data (utils.py and train_cnn.py)
70+
71+
PyTorch provides the DataLoader class to shuffle and split data into batches for training. To use it, we implement a custom Dataset class to read, store, and retrieve our images. The following implementation of `cat_breed_dataset` iterates through every image in our dataset and saves them in an array for images and another array for labels. Here, we assume that in the directory pointed to by `path_to_data`, there are two directories, one holding the images and the other holding the labels, both ordered the same as the other with each image and data called "image#.npy", "#" being a number.
72+
73+
```
74+
class cat_breed_dataset(VisionDataset):
75+
def __init__(self, path_to_data):
76+
self.imgpath = os.path.join(path_to_data, "images")
77+
self.labpath = os.path.join(path_to_data, "labels")
78+
self.images, self.labels = [], []
79+
for lab in os.listdir(self.labpath):
80+
if "item" in lab:
81+
self.labels.append(
82+
np.load(os.path.join(self.labpath, lab))
83+
)
84+
self.labels = np.array(self.labels)
85+
for img in os.listdir(self.imgpath):
86+
if "item" in img:
87+
self.images.append(
88+
np.load(os.path.join(self.imgpath, img))
89+
)
90+
self.images = np.array(self.images)
91+
92+
assert len(self.images) == len(self.labels)
93+
```
94+
95+
Next, we implement `__len__` and `__getitem__`. For the latter, we return the image and the label as two separate tensors. We ensure that the label contains long(s) and the image contains floats for compatibility with the model's weights and the loss function.
96+
97+
```
98+
def __len__(self):
99+
return len(self.labels)
100+
101+
def __getitem__(self, idx):
102+
if torch.is_tensor(idx):
103+
idx = idx.tolist()
104+
img = torch.tensor(self.images[idx]).float()
105+
lab = torch.tensor(self.labels[idx]).long()
106+
# `tensor` is lowercase to make `lab` a 0-dim tensor
107+
return img, lab
108+
```
109+
110+
This class is instantiated and passed to a DataLoader for training.
111+
112+
```
113+
cats_train = utils.cat_breed_dataset("./cat_breeds/data/catbreeds")
114+
trainloader = DataLoader(cats_train, batch_size=128, shuffle=True)
115+
```
116+
117+
## Training the Model (utils.py and train_cnn.py)
118+
119+
Before training, define a loss function and an optimizer to train the model.
120+
121+
```
122+
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
123+
criterion = torch.nn.CrossEntropyLoss()
124+
```
125+
126+
During training, loss is calculated with the model's predictions and the provided labels. The model backpropagates the prediction error for the current batch, adjusting its parameters, before going to the next batch of data until the training is complete.
127+
128+
```
129+
def train(model, device, train_loader, optimizer, criterion):
130+
model.train()
131+
avg_loss = 0
132+
133+
for batch_idx, (data, target) in enumerate(train_loader):
134+
data, target = data.to(device), target.to(device)
135+
optimizer.zero_grad()
136+
output = model(data)
137+
138+
loss = criterion(output, target)
139+
avg_loss += loss
140+
loss.backward()
141+
142+
optimizer.step()
143+
if one_pass: break
144+
145+
avg_loss /= len(train_loader.dataset)
146+
147+
if verbose:
148+
print(f'Average loss: {avg_loss:.6f}')
149+
```
150+
151+
Next, train the model and save it as a .pt file.
152+
153+
for epoch in range(epochs):
154+
utils.train(model, device, trainloader, optimizer, criterion, epoch, one_pass=False, verbose=True)
155+
156+
model.to(torch.device("cpu"))
157+
torch.save(model, "./cat_breeds/models/pretest.pt")
158+
159+
## From PyTorch to ChAI (to_chai.py)
160+
161+
Once the images, labels, and the model have been saved as .npy and .pt files, we can call .chai_dump and .chai_save to save them as files that are readable by the current ChAI functionality. The following saves the first 20 images for brevity.
162+
163+
```
164+
import lib.chai
165+
import torch
166+
import os
167+
import numpy as np
168+
169+
model = torch.load("./cat_breeds/models/pretest.pt")
170+
model.chai_dump("./cat_breeds/models/chai_model", "SmallCNN")
171+
172+
load_path = "./cat_breeds/data/catbreeds/images"
173+
for i, item in enumerate(os.listdir(load_path)):
174+
if "item" in item: # check file name
175+
img = np.load(f"{load_path}/{item}")
176+
img = torch.Tensor(img)
177+
img.chai_save("./cat_breeds/data/catbreeds/chai_images", f"item{i}", verbose=False)
178+
if i > 20:
179+
break
180+
```
181+
182+
The specific path that we follow here holds the data and the model in separate directories, as follows.
183+
184+
cat_breeds
185+
├───models
186+
│ ├───chai_model
187+
│ │ ├───conv1.bias.chdata
188+
│ │ ├───conv2.bias.json
189+
│ │ └───...
190+
│ └───pretest.pt
191+
└───data
192+
└───catbreeds
193+
├───chai_images
194+
│ ├───item0.chdata
195+
│ ├───item0.json
196+
│ └───...
197+
├───images
198+
│ ├───item0.npy
199+
│ ├───item1.npy
200+
│ └───...
201+
└───labels
202+
├───item0.npy
203+
├───item1.npy
204+
└───...
205+
206+
## Single-locale inference in ChAI (single_locale.chpl)
207+
208+
We can call `loadModel` to read the model's information into ChAI.
209+
210+
```
211+
var model: owned Module(real(32)) = loadModel(
212+
specFile="./cat_breeds/models/chai_model/specification.json",
213+
weightsFolder = "./cat_breeds/models/chai_model/",
214+
dtype=real(32)
215+
);
216+
217+
writeln(model.signature);
218+
```
219+
220+
Next, we can call Tensor.load to read each images' data into ChAI. The following code reads `numImages` images into an array.
221+
222+
```
223+
config const numImages = 1;
224+
var images = forall i in 0..<numImages do Tensor.load("./cat_breeds/data/catbreeds/chai_images/item"+i:string+".chdata") : real(32);
225+
```
226+
227+
Lastly, we can use the model by passing images into it, which will call its forward function. The following code passes `numImages` images into the model `numTimes` times.
228+
229+
```
230+
var preds: [0..<numImages] int;
231+
config const numTimes = 1;
232+
var time: real;
233+
for i in 0..<numTimes {
234+
writeln("Inference (loop ",i,")...");
235+
var st = new Time.stopwatch();
236+
237+
st.start();
238+
forall (img, pred) in zip(images, preds) {
239+
writeln(img.type:string);
240+
pred = model(img).argmax();
241+
}
242+
st.stop();
243+
244+
const tm = st.elapsed();
245+
writeln("Time: ", tm, " seconds.");
246+
}
247+
```
Lines changed: 11 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,33 +1,18 @@
1-
import torch
21
import torch.nn as nn
3-
import torch.nn.functional as F
42

53
class SmallCNN(nn.Module):
64
def __init__(self):
75
super(SmallCNN, self).__init__()
8-
self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
9-
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
10-
self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
11-
# self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
12-
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
13-
# self.dropout1 = nn.Dropout(0.25)
14-
self.flat = nn.Flatten()
15-
self.fc1 = nn.Linear(8192, 256)
16-
self.fc2 = nn.Linear(256, 12)
6+
self.layers = nn.Sequential(
7+
nn.Conv2d(3, 64, 3, padding="same"),
8+
nn.ReLU(),
9+
nn.Conv2d(64, 128, 3, padding="same"),
10+
nn.ReLU(),
11+
nn.MaxPool2d(2, 2),
12+
nn.Flatten(),
13+
nn.Linear(8192, 256),
14+
nn.Linear(256, 10)
15+
)
1716

1817
def forward(self, x):
19-
x = F.relu(self.conv1(x))
20-
x = self.pool1(x)
21-
x = F.relu(self.conv2(x))
22-
x = self.pool2(x)
23-
# x = x.view(-1, 8192)
24-
x = self.flat(x)
25-
x = self.fc1(x)
26-
x = self.fc2(x)
27-
# x = F.softmax(x, dim=1)
28-
return x
29-
30-
# NOTES
31-
# Must use nn.Flatten instead of x.view()
32-
# Must use pool1 and pool2 rather than one pool.
33-
# Must not use "same" as padding, because this breaks getInt from moduleAttributes.
18+
return self.layers(x)

examples/cat_breeds/to_chai.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,11 @@
1010
model = torch.load("./cat_breeds/models/pretest.pt")
1111
model.chai_dump("./cat_breeds/models/chai_model", "SmallCNN")
1212

13-
# load_path = "./cat_breeds/data/catbreeds/images"
14-
# for i, item in enumerate(os.listdir(load_path)):
15-
# if "item" in item: # check file name
16-
# img = np.load(f"{load_path}/{item}")
17-
# img = torch.Tensor(img)
18-
# img.chai_save("./cat_breeds/data/catbreeds/chai_images", f"item{i}", verbose=False)
19-
# if i > 20:
20-
# break
13+
load_path = "./cat_breeds/data/catbreeds/images"
14+
for i, item in enumerate(os.listdir(load_path)):
15+
if "item" in item: # check file name
16+
img = np.load(f"{load_path}/{item}")
17+
img = torch.Tensor(img)
18+
img.chai_save("./cat_breeds/data/catbreeds/chai_images", f"item{i}", verbose=False)
19+
if i > 20:
20+
break

lib/Autograd.chpl

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1077,24 +1077,30 @@ record nllLossOp : serializable {
10771077

10781078

10791079
record batchNormOp : serializable {
1080-
type eltType = defaultEltType;
1080+
type eltType = real;
10811081
var features: shared BaseTensorResource(?); // what to put here?
10821082
var weight: shared BaseTensorResource(eltType, 1);
10831083
var bias: shared BaseTensorResource(eltType, 1);
10841084
var movingAvg: shared BaseTensorResource(eltType, 1);
10851085
var movingVar: shared BaseTensorResource(eltType, 1);
1086+
var eps: real;
1087+
var momentum: real;
1088+
var train: bool;
10861089
var n: int;
10871090

1088-
proc children do return (features, weight, bias, movingAvg, movingVar);
1091+
proc children do return (features, weight, bias, movingAvg, movingVar, train);
10891092

10901093
proc forward() {
1091-
return ndarray.batchNorm(features.array, weight.array, bias.array, movingAvg.array, movingVar.array, n);
1094+
if train {
1095+
return ndarray.batchNormTrain(features.array, weight.array, bias.array, movingAvg.array, movingVar.array, eps, momentum, n);
1096+
} else {
1097+
return ndarray.batchNorm(features.array, weight.array, bias.array, movingAvg.array, movingVar.array, eps);
1098+
}
10921099
}
10931100

10941101
proc spec : GradOpSpec do return new dict(("operation","BatchNorm"));
10951102
}
10961103

1097-
10981104
record dropoutOp : serializable {
10991105
param rank: int;
11001106
type eltType;

lib/DynamicTensor.chpl

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -610,9 +610,11 @@ proc type dynamicTensor.batchnorm(
610610
bias: dynamicTensor(eltType),
611611
movingAvg: dynamicTensor(eltType),
612612
movingVar: dynamicTensor(eltType),
613-
numFeatures: int
613+
eps: real,
614+
momentum: real,
615+
train: bool,
616+
num_features: int
614617
): dynamicTensor(eltType) {
615-
616618
for param rankF in 2..4 {
617619
if features.checkRank(rankF) {
618620
return staticTensor.batchNorm(
@@ -621,7 +623,10 @@ proc type dynamicTensor.batchnorm(
621623
bias.forceRank(1),
622624
movingAvg.forceRank(1),
623625
movingVar.forceRank(1),
624-
numFeatures
626+
eps,
627+
momentum,
628+
train,
629+
num_features
625630
).eraseRank();
626631
}
627632
}

0 commit comments

Comments
 (0)