| 174 | } |
| 175 | |
| 176 | TEST(TestSublinearMemory, ConcatSplit) { |
| 177 | CHECK_REQ; |
| 178 | |
| 179 | HostTensorGenerator<> gen_; |
| 180 | auto gen = [&](const TensorShape& shp) { return gen_(shp, "gpu0"); }; |
| 181 | constexpr size_t N = 128, H = 256, W = 256; |
| 182 | auto host_data = gen({N, 2, H, W}); |
| 183 | |
| 184 | auto graph = ComputingGraph::make(); |
| 185 | SymbolVarArray params; |
| 186 | |
| 187 | auto data = opr::Host2DeviceCopy::make(*graph, host_data).rename("data"), |
| 188 | out = data; |
| 189 | size_t out_chl = host_data->shape(1), layer_count = 0; |
| 190 | auto add_layer = [&](size_t oc, size_t h, size_t w) { |
| 191 | auto prev = opr::Split::make(out, opr::Split::Options::make_average(1, 2)); |
| 192 | SymbolVarArray cur_out(2); |
| 193 | size_t cur_in_chl[] = {out_chl / 2, out_chl - out_chl / 2}; |
| 194 | size_t cur_out_chl[] = {oc / 2, oc - oc / 2}; |
| 195 | for (int i = 0; i < 2; ++i) { |
| 196 | gen_.std(sqrt(2.0 / (cur_in_chl[i] * h * w))); |
| 197 | auto host_kern = gen({cur_out_chl[i], cur_in_chl[i], h, w}); |
| 198 | auto dev_kern = std::make_shared<DeviceTensorND>(); |
| 199 | dev_kern->copy_from(*host_kern); |
| 200 | params.emplace_back(opr::SharedDeviceTensor::make(*graph, dev_kern)); |
| 201 | cur_out[i] = opr::relu(opr::Convolution::make( |
| 202 | prev[i], |
| 203 | params.back().rename(ssprintf( |
| 204 | "param%zu:%d", layer_count, i)), |
| 205 | {})) |
| 206 | .rename(ssprintf("out%zu:%d", layer_count, i)); |
| 207 | } |
| 208 | ++layer_count; |
| 209 | out_chl = oc; |
| 210 | out = opr::Concat::make(cur_out, 1); |
| 211 | }; |
| 212 | |
| 213 | for (int i = 0; i < 10; ++i) |
| 214 | add_layer(6, 3, 3); |
| 215 | |
| 216 | auto loss = opr::Dot::make(out.flatten(), out.flatten()); |
| 217 | std::vector<HostTensorND> grad_params_get(params.size()); |
| 218 | ComputingGraph::OutputSpec out_spec; |
| 219 | for (size_t i = 0; i < params.size(); ++i) { |
| 220 | out_spec.emplace_back( |
| 221 | make_callback_copy(cg::grad(loss, params[i]), grad_params_get[i])); |
| 222 | } |
| 223 | |
| 224 | std::vector<HostTensorND> grad_params_expect(grad_params_get.size()); |
| 225 | for (bool sublinear : {false, true}) { |
| 226 | graph->options().enable_sublinear_memory_opt = sublinear; |
| 227 | auto func = graph->compile(out_spec); |
| 228 | func->execute(); |
| 229 | if (!sublinear) { |
| 230 | for (size_t i = 0; i < grad_params_get.size(); ++i) |
| 231 | grad_params_expect[i].copy_from(grad_params_get[i]); |
| 232 | } |
| 233 | } |
no test coverage detected